diff --git a/.gitattributes b/.gitattributes index b11c48ca86039c7cfb465008ca0f286fa44e02d8..0e39929c26bade047f63b72ec308d5c58012ca5f 100644 --- a/.gitattributes +++ b/.gitattributes @@ -36,3 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.jpg filter=lfs diff=lfs merge=lfs -text *.jpeg filter=lfs diff=lfs merge=lfs -text SKING_DDJ_v104/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json filter=lfs diff=lfs merge=lfs -text +SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 963d62d27a973a847e4cb23c8ec2ca03706a69ad..15fb6779cd0e6655660b10529f0fbdc42bb267b6 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,12 @@ tags: **Sking** is an image-to-Minecraft-skin generation model designed to transform character images, portraits, avatars, or concept art into Minecraft-style skin textures. Given a reference image, the model learns to preserve key visual features such as clothing, colors, hairstyle, and character identity while adapting them into the compact pixel-art structure used by Minecraft skins. +## SKING_DDJ_v101c — recommended production UV release (2026-09-16) + +The portable v101c render-to-UV bundle is available in [`SKING_DDJ_v101c/`](SKING_DDJ_v101c). It contains privacy-sanitized inference checkpoints with tensor weights identical to the deployed v101 crown-geometry release, the trained foreground model, pinned offline dependencies, UV mappings, source code, checksums, and a portable launcher. + +See the [v101c download and inference instructions](SKING_DDJ_v101c/README.md). Download the complete directory rather than only `parser.pt`. The older v104 release remains available below for reproducibility. + ## SKING_DDJ_v104 — production UV release (2026-09-08) The full v104 render-to-UV bundle is available in [`SKING_DDJ_v104/`](SKING_DDJ_v104). It includes the dense UV parser, trained foreground model, BiRefNet and SigLIP2 dependencies, pinned runtime, UV mappings, and a portable inference launcher. It produces a 64×64 RGBA Minecraft skin from the combined front/back edited render. diff --git a/SKING_DDJ_v101c/LICENSE b/SKING_DDJ_v101c/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..0ad25db4bd1d86c452db3f9602ccdbe172438f52 --- /dev/null +++ b/SKING_DDJ_v101c/LICENSE @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published + by the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/SKING_DDJ_v101c/README.md b/SKING_DDJ_v101c/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8143c7c24a71696f923be7fd59ee44cd30fb6b1b --- /dev/null +++ b/SKING_DDJ_v101c/README.md @@ -0,0 +1,64 @@ +# SKING_DDJ_v101c + +Portable open-source package of the v101 production render-to-UV pipeline, published as `SKING_DDJ_v101c`. It converts a square combined front/back Minecraft character render into a 64×64 RGBA skin with inner and outer layers. The left view is `front_left` and the right view is `back_left`, following the Sking edited-render layout. Creating that edited render is a separate upstream stage. + +The parser tensor weights, foreground decoder tensor weights, pipeline behavior, and runtime are derived from the deployed v101 crown-geometry release (`crown_geometry_20260906`, source commit `28f548a5be37868022cfbafaddbe095895e18210`). The public checkpoints were repacked to remove optimizer state, private filesystem paths, and development-only manifests. Tensor equality and inference outputs were verified against the deployed release. + +## Contents + +| Path | Purpose | +| --- | --- | +| `parser.pt` | Sanitized v101c dense UV parser checkpoint | +| `foreground.pt` | Sanitized trained Minecraft foreground decoder checkpoint | +| `pipeline.json` | Inference configuration with development-only path metadata removed | +| `foreground_base/` | BiRefNet base model and implementation | +| `hf_cache/` | Pinned SigLIP2 model and tokenizer for offline inference | +| `runtime/` | Matching inference and training source snapshot, without private test images | +| `mappings_256x512/` | UV mapping tensors used by production | +| `infer.py` | Portable foreground + parser + final inpainting launcher | +| `release.json` | Checksums, source artifact identities, and sanitization record | +| `tested_environments.json` | Dependency versions used for verification | + +Download the complete directory. `parser.pt` by itself is insufficient to reproduce the full pipeline. + +## Download + +```python +from huggingface_hub import snapshot_download + +snapshot_download( + repo_id="EntropyDrop/Sking", + allow_patterns=["SKING_DDJ_v101c/*"], + local_dir="./Sking", +) +``` + +For repeatable deployments, pass the release commit hash as `revision` instead of following the repository's moving `main` branch. + +## Run + +Use Linux with an NVIDIA CUDA GPU and Python 3.12. The parser environment needs PyTorch/torchvision, NumPy, Pillow, Transformers, timm, einops, SciPy, OpenCV, SentencePiece, safetensors and accelerate. The foreground environment also uses kornia. Install a CUDA-compatible PyTorch build before the remaining dependencies. The two options below may point to the same interpreter if it contains both dependency sets. + +```bash +python ./Sking/SKING_DDJ_v101c/infer.py --verify-only + +/path/to/parser-env/bin/python ./Sking/SKING_DDJ_v101c/infer.py \ + --parser-python /path/to/parser-env/bin/python \ + --foreground-python /path/to/foreground-env/bin/python \ + --input ./character_edited.png \ + --output ./character_result_v101c.png +``` + +Input must be an even square image between 256 and 4096 pixels wide, containing the two expected views side by side. The launcher verifies every bundled file, runs the trained foreground model, executes the pinned parser and postprocessing pipeline, validates output provenance, and writes the final 64×64 RGBA PNG. Model inference uses the included Hugging Face cache in offline mode. + +## Validation and limits + +Eight fresh samples were selected across the historical v54, v61, and v66 data groups, excluding the deployment fixtures. Public-package outputs had exact alpha on all eight samples and visible RGB differences of at most 1/255 from the historical v101 results. A separate six-case deployment suite covering eyes, hair, beard, crown, hat, and glasses met the same threshold. + +The repacked parser and foreground checkpoints were checked tensor-by-tensor against production. The public package contains no optimizer states, credentials, server paths, user inputs, deployment records, backup directories, or historical regression image galleries. + +Semantic errors can still occur, and the model assumes the expected edited-render layout. + +## License and upstream components + +Sking is licensed under AGPL-3.0; see `LICENSE`. BiRefNet and SigLIP2 retain their own licenses and attribution. See `THIRD_PARTY_NOTICES.md` and `licenses/` for pinned upstream revisions and license texts. diff --git a/SKING_DDJ_v101c/THIRD_PARTY_NOTICES.md b/SKING_DDJ_v101c/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000000000000000000000000000000000000..78b2b0087977b8f8be8c6baa45f61ff849988793 --- /dev/null +++ b/SKING_DDJ_v101c/THIRD_PARTY_NOTICES.md @@ -0,0 +1,7 @@ +# Third-party components + +- `foreground_base/`: [ZhengPeng7/BiRefNet](https://huggingface.co/ZhengPeng7/BiRefNet/tree/6a62b7dcfa18a3829087877fb16c8006831e4220), MIT. License: `licenses/BiRefNet_LICENSE.txt`. +- `hf_cache/hub/models--google--siglip2-base-patch16-224/`: [google/siglip2-base-patch16-224](https://huggingface.co/google/siglip2-base-patch16-224/tree/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2), Apache-2.0. License: `licenses/SigLIP2_LICENSE.txt` (Google Big Vision repository). +- Sking weights, inference/runtime code and UV mappings are provided under the repository AGPL-3.0 license; third-party components retain their respective licenses. + +Upstream license retrieval URLs and SHA-256 hashes are recorded in `licenses/sources.json`. diff --git a/SKING_DDJ_v101c/foreground.pt b/SKING_DDJ_v101c/foreground.pt new file mode 100644 index 0000000000000000000000000000000000000000..ae9b3f1c23798e8fab7fdcb9ff76f61130c1c909 --- /dev/null +++ b/SKING_DDJ_v101c/foreground.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f00c1abbf736d54fcda6800ec6046817456ef9fbc45bb67bc1e6d759cfb45331 +size 100142558 diff --git a/SKING_DDJ_v101c/foreground_base/BiRefNet_config.py b/SKING_DDJ_v101c/foreground_base/BiRefNet_config.py new file mode 100644 index 0000000000000000000000000000000000000000..37c8ac58bec2f52dac34204978a7b61b69e3da76 --- /dev/null +++ b/SKING_DDJ_v101c/foreground_base/BiRefNet_config.py @@ -0,0 +1,11 @@ +from transformers import PretrainedConfig + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + def __init__( + self, + bb_pretrained=False, + **kwargs + ): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/SKING_DDJ_v101c/foreground_base/birefnet.py b/SKING_DDJ_v101c/foreground_base/birefnet.py new file mode 100644 index 0000000000000000000000000000000000000000..57eba49a236a93365582e59814641e01f4facaa8 --- /dev/null +++ b/SKING_DDJ_v101c/foreground_base/birefnet.py @@ -0,0 +1,2249 @@ +### config.py + +import os +import math +from transformers import PretrainedConfig + + +class Config(PretrainedConfig): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from .BiRefNet_config import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/SKING_DDJ_v101c/foreground_base/config.json b/SKING_DDJ_v101c/foreground_base/config.json new file mode 100644 index 0000000000000000000000000000000000000000..06d8fa9d7f2f4c6f1cf0dc6e7bfd194153176a42 --- /dev/null +++ b/SKING_DDJ_v101c/foreground_base/config.json @@ -0,0 +1,20 @@ +{ + "_name_or_path": "ZhengPeng7/BiRefNet", + "architectures": [ + "BiRefNet" + ], + "auto_map": { + "AutoConfig": "BiRefNet_config.BiRefNetConfig", + "AutoModelForImageSegmentation": "birefnet.BiRefNet" + }, + "custom_pipelines": { + "image-segmentation": { + "pt": [ + "AutoModelForImageSegmentation" + ], + "tf": [], + "type": "image" + } + }, + "bb_pretrained": false +} \ No newline at end of file diff --git a/SKING_DDJ_v101c/foreground_base/model.safetensors b/SKING_DDJ_v101c/foreground_base/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..3b958f6368413c833b0abc357a1357cb78ce6d52 --- /dev/null +++ b/SKING_DDJ_v101c/foreground_base/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9ab37426bf4de0567af6b5d21b16151357149139362e6e8992021b8ce356a154 +size 444473596 diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/refs/main b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/refs/main new file mode 100644 index 0000000000000000000000000000000000000000..5241bfe843c9dd8a7c06a84caa2b3787385e3c0f --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/refs/main @@ -0,0 +1 @@ +75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2 \ No newline at end of file diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/config.json b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/config.json new file mode 100644 index 0000000000000000000000000000000000000000..c8cd2a20e58a738f44f267ae19f9568ff1095698 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/config.json @@ -0,0 +1,12 @@ +{ + "initializer_factor": 1.0, + "model_type": "siglip", + "text_config": { + "model_type": "siglip_text_model", + "vocab_size": 256000 + }, + "transformers_version": "4.49.0.dev0", + "vision_config": { + "model_type": "siglip_vision_model" + } +} diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/model.safetensors b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..aff60bf4990ccd2fc9de54afb14d8bd3580cdf8a --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:612923381c76ec5a9bed335d1c48827e3f2e506ac31b044b63b2031fadee6a0b +size 1500800904 diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/preprocessor_config.json b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..2e52d8e8492b5c496ae04c37bfa09760469fb18b --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/preprocessor_config.json @@ -0,0 +1,24 @@ +{ + "do_convert_rgb": null, + "do_normalize": true, + "do_rescale": true, + "do_resize": true, + "image_mean": [ + 0.5, + 0.5, + 0.5 + ], + "image_processor_type": "SiglipImageProcessor", + "image_std": [ + 0.5, + 0.5, + 0.5 + ], + "processor_class": "SiglipProcessor", + "resample": 2, + "rescale_factor": 0.00392156862745098, + "size": { + "height": 224, + "width": 224 + } +} diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/special_tokens_map.json b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/special_tokens_map.json new file mode 100644 index 0000000000000000000000000000000000000000..8d6368f7e735fbe4781bf6e956b7c6ad0586df80 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/special_tokens_map.json @@ -0,0 +1,34 @@ +{ + "additional_special_tokens": [ + "", + "" + ], + "bos_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "eos_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "pad_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + }, + "unk_token": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false + } +} diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json new file mode 100644 index 0000000000000000000000000000000000000000..16d02f7b98cf1fe1ee9bf04143543b2f188f2788 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cb9140fae3ac5122c972d37adf83e1248471a38147ad76f8215c8872c6fd8322 +size 34363039 diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.model b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.model new file mode 100644 index 0000000000000000000000000000000000000000..796efe9ab515c15e146ce7588e6d7b9b8134dbf8 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.model @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2 +size 4241003 diff --git a/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer_config.json b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer_config.json new file mode 100644 index 0000000000000000000000000000000000000000..d97c5412159422c3b56fbc99076b1dcc25dd7856 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer_config.json @@ -0,0 +1,2020 @@ +{ + "add_bos_token": false, + "add_eos_token": true, + "added_tokens_decoder": { + "0": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "1": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "2": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "3": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "4": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "5": { + "content": "<2mass>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "6": { + "content": "[@BOS@]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "7": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "8": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "9": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "10": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "11": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "12": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "13": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "14": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "15": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "16": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "17": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "18": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "19": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "20": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "21": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "22": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "23": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "24": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "25": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "26": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "27": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "28": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "29": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "30": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "31": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "32": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "33": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "34": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "35": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "36": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "37": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "38": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "39": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "40": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "41": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "42": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "43": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "44": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "45": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "46": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "47": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "48": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "49": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "50": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "51": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "52": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "53": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "54": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "55": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "56": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "57": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "58": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "59": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "60": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "61": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "62": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "63": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "64": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "65": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "66": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "67": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "68": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "69": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "70": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "71": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "72": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "73": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "74": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "75": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "76": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "77": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "78": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "79": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "80": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "81": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "82": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "83": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "84": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "85": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "86": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "87": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "88": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "89": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "90": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "91": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "92": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "93": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "94": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "95": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "96": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "97": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "98": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "99": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "100": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "101": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "102": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "103": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "104": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "105": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "106": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "107": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "108": { + "content": "\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "109": { + "content": "\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "110": { + "content": "\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "111": { + "content": "\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "112": { + "content": "\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "113": { + "content": "\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "114": { + "content": "\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "115": { + "content": "\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "116": { + "content": "\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "117": { + "content": "\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "118": { + "content": "\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "119": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "120": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "121": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "122": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "123": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "124": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "125": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "126": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "127": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "128": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "129": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "130": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "131": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "132": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "133": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "134": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "135": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "136": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "137": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "138": { + "content": "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "139": { + "content": "▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "140": { + "content": "▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "141": { + "content": "▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "142": { + "content": "▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "143": { + "content": "▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "144": { + "content": "▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "145": { + "content": "▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "146": { + "content": "▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "147": { + "content": "▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "148": { + "content": "▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "149": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "150": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "151": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "152": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "153": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "154": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "155": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "156": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "157": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "158": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "159": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "160": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "161": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "162": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "163": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "164": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "165": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "166": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "167": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "168": { + "content": "▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "169": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "170": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "172": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "173": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "174": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "175": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "171": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "176": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "177": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "178": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "179": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "180": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "181": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "182": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "183": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "184": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "185": { + "content": "

", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "186": { + "content": "

", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "187": { + "content": "

", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "188": { + "content": "

", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "189": { + "content": "

", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "190": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "191": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "192": { + "content": "
", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "193": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "194": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "195": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "196": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "197": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "198": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "199": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "200": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "201": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "202": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "203": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "204": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "205": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "206": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "207": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "208": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "209": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "210": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "211": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "212": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "213": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "214": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "215": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "216": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255968": { + "content": "[toxicity=0]", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255969": { + "content": "\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255970": { + "content": "\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255971": { + "content": "\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255972": { + "content": "\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255973": { + "content": "\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255974": { + "content": "\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255975": { + "content": "\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255976": { + "content": "\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255977": { + "content": "\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255978": { + "content": "\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255979": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255980": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255981": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255982": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255983": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255984": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255985": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255986": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255987": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255988": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255989": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255990": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255991": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255992": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255993": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255994": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255995": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255996": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255997": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255998": { + "content": "\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + }, + "255999": { + "content": "", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": false + } + }, + "additional_special_tokens": [ + "", + "" + ], + "bos_token": "", + "chat_template": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '' + role + '\n' + message['content'] | trim + '\n' }}{% endfor %}{% if add_generation_prompt %}{{'model\n'}}{% endif %}", + "clean_up_tokenization_spaces": false, + "do_lower_case": true, + "eos_token": "", + "extra_special_tokens": {}, + "model_input_names": [ + "input_ids" + ], + "model_max_length": 1000000000000000019884624838656, + "pad_token": "", + "padding_side": "right", + "processor_class": "SiglipProcessor", + "sp_model_kwargs": {}, + "spaces_between_special_tokens": false, + "tokenizer_class": "GemmaTokenizer", + "unk_token": "", + "use_default_system_prompt": false +} diff --git a/SKING_DDJ_v101c/hf_cache/modules/__init__.py b/SKING_DDJ_v101c/hf_cache/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/__init__.py b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/BiRefNet_config.py b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/BiRefNet_config.py new file mode 100644 index 0000000000000000000000000000000000000000..37c8ac58bec2f52dac34204978a7b61b69e3da76 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/BiRefNet_config.py @@ -0,0 +1,11 @@ +from transformers import PretrainedConfig + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + def __init__( + self, + bb_pretrained=False, + **kwargs + ): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/__init__.py b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/birefnet.py b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/birefnet.py new file mode 100644 index 0000000000000000000000000000000000000000..57eba49a236a93365582e59814641e01f4facaa8 --- /dev/null +++ b/SKING_DDJ_v101c/hf_cache/modules/transformers_modules/foreground_base/birefnet.py @@ -0,0 +1,2249 @@ +### config.py + +import os +import math +from transformers import PretrainedConfig + + +class Config(PretrainedConfig): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from .BiRefNet_config import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/SKING_DDJ_v101c/infer.py b/SKING_DDJ_v101c/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..f32bacdc1ef0a414fc185b8249cf8bf4382321d8 --- /dev/null +++ b/SKING_DDJ_v101c/infer.py @@ -0,0 +1,214 @@ +"""Portable launcher for the sanitized SKING_DDJ_v101c inference bundle.""" +import argparse +import hashlib +import json +import os +from pathlib import Path +import runpy +import shutil +import subprocess +import sys +import tempfile + +ROOT = Path(__file__).resolve().parent + + +def sha(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(4 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def verify(): + release = json.loads((ROOT / "release.json").read_text()) + if ( + release["format"] != "sking-hf-bundle-v1" + or release["model_version"] != "SKING_DDJ_v101c" + ): + raise ValueError("Wrong release identity") + for name, expected in release["sha256"].items(): + path = (ROOT / name).resolve() + if not path.is_relative_to(ROOT) or sha(path) != expected: + raise ValueError("Release integrity failure: " + name) + return release + + +def parser_stage(options): + sys.path.insert(0, str(ROOT / "runtime/dense_uv_parser")) + from run_local import bind_checkout + + bind_checkout() + import SkingToolkit.dense_uv_parser.infer as parser_module + + original = parser_module.load_parser + + def load_portable(*args, **kwargs): + model, config = original(*args, **kwargs) + config = dict(config) + config["mappings_dir"] = str(ROOT / "mappings_256x512") + return model, config + + parser_module.load_parser = load_portable + sys.argv = [ + "batch_accessories", + "--checkpoint", + str(ROOT / "parser.pt"), + "--pipeline", + str(ROOT / "pipeline.json"), + "--output-dir", + str(options.output), + "--foreground-probabilities", + str(options.probability), + "--inputs", + str(options.input), + ] + runpy.run_module( + "SkingToolkit.dense_uv_parser.batch_accessories", run_name="__main__" + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, help="Square combined front/back edited render") + parser.add_argument("--output", type=Path, help="Output 64x64 RGBA PNG") + parser.add_argument("--parser-python", default=sys.executable) + parser.add_argument("--foreground-python", default=sys.executable) + parser.add_argument("--verify-only", action="store_true") + parser.add_argument("--parser-stage", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--probability", type=Path, help=argparse.SUPPRESS) + options = parser.parse_args() + os.environ.update( + HF_HUB_OFFLINE="1", + HF_HOME=str(ROOT / "hf_cache"), + HF_HUB_CACHE=str(ROOT / "hf_cache/hub"), + HF_HUB_DISABLE_PROGRESS_BARS="1", + PYTHONDONTWRITEBYTECODE="1", + OMP_NUM_THREADS="4", + MKL_NUM_THREADS="4", + CUBLAS_WORKSPACE_CONFIG=":4096:8", + ) + if options.parser_stage: + parser_stage(options) + return + release = verify() + if options.verify_only: + print( + json.dumps( + { + "verified_files": len(release["sha256"]), + "model_version": release["model_version"], + } + ) + ) + return + if options.input is None or options.output is None: + parser.error("--input and --output are required for inference") + options.input = options.input.resolve() + options.output = options.output.resolve() + if options.input == options.output: + parser.error("Input and output must be different files") + from PIL import Image + + with Image.open(options.input) as source: + width, height = source.size + if width != height or width < 256 or width > 4096 or width % 2: + raise ValueError("Expected an even square combined render, 256 to 4096 pixels") + source.verify() + options.output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="sking-v101c-") as directory: + work = Path(directory) + source_path = work / "input.png" + shutil.copyfile(options.input, source_path) + index = work / "foreground.json" + entry = ROOT / "runtime/dense_uv_parser/run_local.py" + subprocess.run( + [ + options.foreground_python, + str(entry), + "foreground_provider", + "--checkpoint", + str(ROOT / "foreground.pt"), + "--model-dir", + str(ROOT / "foreground_base"), + "--output-dir", + str(work / "foreground"), + "--result-manifest", + str(index), + "--inputs", + str(source_path), + ], + check=True, + cwd=ROOT / "runtime", + ) + probabilities = json.loads(index.read_text())["probabilities"] + if len(probabilities) != 1: + raise RuntimeError("Foreground stage did not return one result") + subprocess.run( + [ + options.parser_python, + str(Path(__file__).resolve()), + "--parser-stage", + "--input", + str(source_path), + "--output", + str(work / "uv"), + "--probability", + probabilities[0], + ], + check=True, + cwd=ROOT / "runtime", + ) + outputs = list((work / "uv").glob("*/manifest.json")) + if len(outputs) != 1: + raise RuntimeError("Parser did not return one result") + manifest = json.loads(outputs[0].read_text()) + expected_sources = { + Path(name).name: digest + for name, digest in release["sha256"].items() + if Path(name).parent == Path("runtime/dense_uv_parser") + and name.endswith(".py") + } + if ( + manifest.get("complete") is not True + or manifest.get("input_sha256") != sha(source_path) + or manifest.get("checkpoint_sha256") != release["sha256"]["parser.pt"] + or manifest.get("pipeline") + != json.loads((ROOT / "pipeline.json").read_text()) + or manifest.get("source_sha256") != expected_sources + or manifest.get("foreground_provider", {}).get("adaptation_sha256") + != release["sha256"]["foreground.pt"] + or manifest.get("foreground_provider", {}).get("model_sha256") + != release["sha256"]["foreground_base/model.safetensors"] + ): + raise RuntimeError("Output provenance does not match this release") + result = outputs[0].parent / "pred_uv.png" + with Image.open(result) as uv: + if uv.size != (64, 64) or uv.mode != "RGBA" or not uv.getchannel("A").getbbox(): + raise RuntimeError("Invalid UV output") + with tempfile.NamedTemporaryFile( + dir=options.output.parent, + prefix=".sking-v101c-", + suffix=".png", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(result.read_bytes()) + try: + os.replace(temporary, options.output) + finally: + temporary.unlink(missing_ok=True) + print( + json.dumps( + { + "model_version": "SKING_DDJ_v101c", + "output": str(options.output), + "sha256": sha(options.output), + } + ) + ) + + +if __name__ == "__main__": + main() diff --git a/SKING_DDJ_v101c/licenses/BiRefNet_LICENSE.txt b/SKING_DDJ_v101c/licenses/BiRefNet_LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..485921e8271834871ff2baafb892b79ef6a75e75 --- /dev/null +++ b/SKING_DDJ_v101c/licenses/BiRefNet_LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 ZhengPeng + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/SKING_DDJ_v101c/licenses/SigLIP2_LICENSE.txt b/SKING_DDJ_v101c/licenses/SigLIP2_LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..f49a4e16e68b128803cc2dcea614603632b04eac --- /dev/null +++ b/SKING_DDJ_v101c/licenses/SigLIP2_LICENSE.txt @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/SKING_DDJ_v101c/licenses/sources.json b/SKING_DDJ_v101c/licenses/sources.json new file mode 100644 index 0000000000000000000000000000000000000000..cb17b943e49189c40a7b8f2170d6f1cdcbb79e9e --- /dev/null +++ b/SKING_DDJ_v101c/licenses/sources.json @@ -0,0 +1,14 @@ +[ + { + "file": "BiRefNet_LICENSE.txt", + "url": "https://raw.githubusercontent.com/ZhengPeng7/BiRefNet/main/LICENSE", + "sha256": "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb", + "bytes": 1066 + }, + { + "file": "SigLIP2_LICENSE.txt", + "url": "https://raw.githubusercontent.com/google-research/big_vision/main/LICENSE", + "sha256": "43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1", + "bytes": 11356 + } +] \ No newline at end of file diff --git a/SKING_DDJ_v101c/mappings_256x512/back_left_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/back_left_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..ccfe7a08b239095d2813e9801467380df645d969 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/back_left_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4966dc30a35f7884dad1f37d4cf2b14bcf8fdfdbe34f3dbbec1ed2ebd31f9308 +size 18878693 diff --git a/SKING_DDJ_v101c/mappings_256x512/back_left_overlay_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/back_left_overlay_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..66d61c2fe04d278aa46133d88f091480066c3ef1 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/back_left_overlay_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:66796f7e6f321f9bd1739fbd9f2110c6477d814cf9a98f968190a255b2e128e7 +size 13373861 diff --git a/SKING_DDJ_v101c/mappings_256x512/back_right_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/back_right_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..9195d1f7c026bab8e9fb1e46c3304b4f0b8b9025 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/back_right_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3baa16ca38e7b62f38d04c29746c72c5a3c05ca054d72300cda2c08b2645c4c8 +size 18878709 diff --git a/SKING_DDJ_v101c/mappings_256x512/back_right_overlay_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/back_right_overlay_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..2378d9575db8e488abbb1e8d19021173efa01212 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/back_right_overlay_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63e7dde6cfb22edebe4b45b1ebecaa9ee5679aca1be022e26fabdfc922c84fb8 +size 13373877 diff --git a/SKING_DDJ_v101c/mappings_256x512/front_left_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/front_left_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..51151e7e06a54f3a9ea1c14743b4c88f8cc97a6d --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/front_left_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f09e38516ed10a81c86939119d776353d795554093c699b10d89741ea78dce7f +size 18878709 diff --git a/SKING_DDJ_v101c/mappings_256x512/front_left_overlay_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/front_left_overlay_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..e2ccc26df38e9d5259b6eb0471b44a6573710497 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/front_left_overlay_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:56db06fe4bf215e06b6f9e8a7b81747e4e24483d770619660b26f555e5402f97 +size 13373877 diff --git a/SKING_DDJ_v101c/mappings_256x512/front_right_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/front_right_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..63accdc319bee4eac98906d6a0bae4668b6bb886 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/front_right_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d9f32bdcbbd98ea6f75aaacc966d6553582d676c06154003dde1559d3f92f1f5 +size 18878725 diff --git a/SKING_DDJ_v101c/mappings_256x512/front_right_overlay_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/front_right_overlay_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..faa9f423c0cc6d9cb4edc3dd02a25ef9fa385b2a --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/front_right_overlay_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:138481fd6da1746f21dd90bb42f432a532fa79d73d60515d4558bda2cf56365d +size 13373893 diff --git a/SKING_DDJ_v101c/mappings_256x512/static_back_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/static_back_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..518bb2eeda5503f4f0970314f8c99ae9d7c6bc49 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/static_back_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14213de89eaac720fd83892033c5e42301e682ad37c678dd723619bf10982ca2 +size 18878725 diff --git a/SKING_DDJ_v101c/mappings_256x512/static_front_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/static_front_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..0a1ffe3f27ea396a44612dcbccd3c3d30d372cf5 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/static_front_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d42e32902fedb9af2bd0f34ae50d18107c3d2264e519beb3b4eb8619d8a0f703 +size 18878741 diff --git a/SKING_DDJ_v101c/mappings_256x512/walk_back_base_layer_ortho_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/walk_back_base_layer_ortho_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..4c4bc44507c128b5f518e96de4dfe661e2fb879f --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/walk_back_base_layer_ortho_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9fdc06695caada150b071bc0632fb83a8fed04c6ae1135b25c9519077eddfdca +size 13374005 diff --git a/SKING_DDJ_v101c/mappings_256x512/walk_back_both_layer_ortho_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/walk_back_both_layer_ortho_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..d6877f279292df75bd2ef3b826240c373ab0859a --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/walk_back_both_layer_ortho_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dfe71bab556365f9e4a33c2eca9d1905806b4278d422adfcf58cb6678b65853 +size 21501093 diff --git a/SKING_DDJ_v101c/mappings_256x512/walk_front_base_layer_ortho_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/walk_front_base_layer_ortho_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..668e346961f468e1cfd510da463d8b4a493f131a --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/walk_front_base_layer_ortho_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c406299c5bab60cc4b8d7afa59f833a41a693ba7b2efa7a819ee2828e4f4008b +size 13374021 diff --git a/SKING_DDJ_v101c/mappings_256x512/walk_front_both_layer_ortho_mapping.pt b/SKING_DDJ_v101c/mappings_256x512/walk_front_both_layer_ortho_mapping.pt new file mode 100644 index 0000000000000000000000000000000000000000..062fe3b64690d6b40975956376ffb8303de3eb57 --- /dev/null +++ b/SKING_DDJ_v101c/mappings_256x512/walk_front_both_layer_ortho_mapping.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:84ad02344ba9777c18d0ec41004d6fed5211b8a1e5cffd62dcac43ba7e4f3148 +size 21501111 diff --git a/SKING_DDJ_v101c/parser.pt b/SKING_DDJ_v101c/parser.pt new file mode 100644 index 0000000000000000000000000000000000000000..d90e3e1c42319ad13df348f5db099f31a30bbf9a --- /dev/null +++ b/SKING_DDJ_v101c/parser.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea3bd7d0e0d2eeb73d0bea68233412bf74ddc5ae8ddfb9a3a6957c381f1b2540 +size 31771763 diff --git a/SKING_DDJ_v101c/pipeline.json b/SKING_DDJ_v101c/pipeline.json new file mode 100644 index 0000000000000000000000000000000000000000..4a7e30f2a426898802f8b1d3f6a8705002c62811 --- /dev/null +++ b/SKING_DDJ_v101c/pipeline.json @@ -0,0 +1,79 @@ +{ + "foreground_flood_tolerance": 0.09, + "foreground_background": "adaptive", + "routing": { + "views": [ + "front_left", + "back_left" + ], + "group_size": 2, + "fg_threshold": 0.5, + "bg_color": [ + 128, + 128, + 128 + ], + "semantic_gate": true, + "affine_refine": false, + "affine_refine_translation_px": 0.0, + "affine_refine_scale": 0.0, + "route_confidence_threshold": 0.0, + "route_margin_threshold": 0.0, + "outer_route_confidence_threshold": 0.8, + "outer_route_margin_threshold": 0.55, + "outer_uv_min_coverage": 0.25, + "outer_uv_min_source_pixels": 30, + "outer_silhouette_consistency": true, + "outer_silhouette_min_coverage": 0.5, + "outer_silhouette_dilation": 0, + "outer_silhouette_min_pixels": 4, + "outer_geometry_rescue": true, + "outer_semantic_rescue": true, + "outer_semantic_presence_threshold": 0.8, + "outer_semantic_coverage_threshold": 0.2, + "outer_rescue_confidence_threshold": 0.6, + "outer_rescue_margin_threshold": 0.25, + "outer_rescue_min_coverage": 0.1, + "color_aggregation": "grid_mode", + "geometry_route_texel_consensus": true, + "geometry_route_texel_consensus_weight": 0.6, + "geometry_route_preserve_outer_confidence": 0.8, + "geometry_route_preserve_outer_margin": 0.35, + "geometry_route_consensus_outer_confidence": 0.7, + "geometry_route_consensus_outer_margin": 0.2, + "geometry_cross_view_outer_consistency": false, + "geometry_cross_view_outer_weight": 0.5, + "geometry_cross_view_outer_positive_confidence": 0.7, + "geometry_cross_view_outer_positive_margin": 0.2, + "geometry_cross_view_outer_negative_confidence": 0.7, + "geometry_cross_view_outer_negative_margin": 0.2, + "geometry_cross_view_outer_background_max_coverage": 0.25, + "geometry_cross_view_outer_min_views": 2, + "outer_uv_occupancy": false, + "outer_uv_occupancy_blend_weight": 0.0, + "outer_uv_occupancy_gate_threshold": 0.15, + "outer_uv_occupancy_rescue_threshold": 0.7, + "outer_uv_occupancy_rescue_route_threshold": 0.3, + "outer_uv_component_routing": false, + "outer_uv_component_seed_threshold": 0.8, + "outer_uv_component_grow_threshold": 0.5, + "outer_uv_component_min_size": 2, + "background_color_tolerance": 0.25, + "color_background_tolerance": 0.031372549, + "color_foreground_inset": 1, + "reject_semantic_fallback": true, + "include_rejected_context": false, + "include_confidence": false, + "head_color_separation_inset": 2 + }, + "accessory_route_threshold": 0.9, + "synthetic_alpha_foreground": true, + "head_material_refine_steps": 48, + "head_material_refine_mode": "isolated", + "headphone_routing_mode": "existing_uv", + "headphone_presence_threshold": 0.99, + "headphone_presence_consensus": true, + "head_semantic_inpaint": true, + "headwear_presence_threshold": 0.95, + "crown_top_geometry_mode": "rendered_semantics" +} diff --git a/SKING_DDJ_v101c/release.json b/SKING_DDJ_v101c/release.json new file mode 100644 index 0000000000000000000000000000000000000000..2d219dab68bd793a14ea4983a82bacc10643a4f5 --- /dev/null +++ b/SKING_DDJ_v101c/release.json @@ -0,0 +1,142 @@ +{ + "format": "sking-hf-bundle-v1", + "model_version": "SKING_DDJ_v101c", + "original_version": "v101", + "original_revision": "crown_geometry_20260906", + "source_code_commit": "28f548a5be37868022cfbafaddbe095895e18210", + "production_artifacts": { + "parser_sha256": "a8aa3d8cd51cc6fec28d7c525aa1b012976d7cb1206e8b00c707de76bfb205dc", + "foreground_sha256": "8585a698d38969e6e7562e9c7d8c5f0b5e934e8f0811cf5b8678fcf05734e732", + "pipeline_sha256": "3e14b98437eba9c5189a3344f560e7183f8daad44b9c40c55e15f7e0fda27685" + }, + "sanitization": { + "pipeline_fields_removed": [ + "foreground_cache", + "foreground_release", + "hat_development" + ], + "private_paths_removed": true, + "optimizer_states_removed": true, + "user_and_regression_images_excluded": true, + "parser_tensor_count": 386, + "foreground_tensor_count": 397, + "parser_metadata_paths_removed": 1, + "foreground_metadata_paths_removed": 3, + "optimizer_state_removed": { + "parser": false, + "foreground": true + }, + "excluded_private_tests": [ + "test_accessories.py", + "test_affine_routing.py", + "test_crown_geometry.py", + "test_foreground.py", + "test_headwear.py", + "test_matting_data.py", + "test_ownership.py", + "test_semantic_generalization.py", + "test_semantic_parser.py" + ], + "portable_path_adjustments": [ + "batch_accessories.py", + "foreground_batch.py", + "released_inference.py" + ] + }, + "validation": { + "historical_samples": 8, + "alpha_exact": 8, + "visible_rgb_max": 1, + "deployment_regression_samples": 6 + }, + "sha256": { + "LICENSE": "8486a10c4393cee1c25392769ddd3b2d6c242d6ec7928e1414efff7dfb2f07ef", + "README.md": "342c7e13ea8a7cad51d77bf999074168ea99db7a48e2c4cc6b7b78e80b362667", + "THIRD_PARTY_NOTICES.md": "58533726d4b29528a102039a8d1c6d42cf6c7b07b42dd9b026cf418348fc14af", + "foreground.pt": "f00c1abbf736d54fcda6800ec6046817456ef9fbc45bb67bc1e6d759cfb45331", + "foreground_base/BiRefNet_config.py": "e7b8c2a74f6cea6a59553d517f71d47f2c1d90e670a13416af17c25fe2f3dc52", + "foreground_base/birefnet.py": "a9566611aa07a6fbb68ddb6ac8e19e62c879a428fbfc2f44efa53d233f2e302f", + "foreground_base/config.json": "c97ea21569daf66b205491a4635147dd3bc42c7c168b89d7d75b53f67ef548ae", + "foreground_base/model.safetensors": "9ab37426bf4de0567af6b5d21b16151357149139362e6e8992021b8ce356a154", + "hf_cache/hub/models--google--siglip2-base-patch16-224/refs/main": "5adaedd88ee9835df44ae379db76ca2e0a84427d427f05d9b268f5556b02ff97", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/config.json": "fe8b5fe6d5734360678fd71c11c21e1ea3364bd8598d34295d9206335973ffd7", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/model.safetensors": "612923381c76ec5a9bed335d1c48827e3f2e506ac31b044b63b2031fadee6a0b", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/preprocessor_config.json": "9b36b57ebaf20f09bf4c22100ccc21877ea6bfe5aead0c00c59f8af8ccefacfc", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/special_tokens_map.json": "baec30ea10906f16adb8c18af7a34023002c1746542612b8b41c9f09e1351351", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.json": "cb9140fae3ac5122c972d37adf83e1248471a38147ad76f8215c8872c6fd8322", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer.model": "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2", + "hf_cache/hub/models--google--siglip2-base-patch16-224/snapshots/75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2/tokenizer_config.json": "14afe629fe4959b9e0d51e1852b8d9f7ad074f90a1a7125a4fcdd17f06e78fc8", + "hf_cache/modules/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "hf_cache/modules/transformers_modules/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "hf_cache/modules/transformers_modules/foreground_base/BiRefNet_config.py": "e7b8c2a74f6cea6a59553d517f71d47f2c1d90e670a13416af17c25fe2f3dc52", + "hf_cache/modules/transformers_modules/foreground_base/__init__.py": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "hf_cache/modules/transformers_modules/foreground_base/birefnet.py": "a9566611aa07a6fbb68ddb6ac8e19e62c879a428fbfc2f44efa53d233f2e302f", + "infer.py": "4454b60fe21440cd8036cb7fdac54063d4b82a1636347c650475a59dcda45c29", + "licenses/BiRefNet_LICENSE.txt": "92a7089e0915fc32bc40067560b398f1e6a7a5958abd7d04eda393629a5acefb", + "licenses/SigLIP2_LICENSE.txt": "43070e2d4e532684de521b885f385d0841030efa2b1a20bafb76133a5e1379c1", + "licenses/sources.json": "61b376b10e8f5d18d7d99c0cb9db89f761e60d1b2ad5333b8408502d535c4218", + "mappings_256x512/back_left_mapping.pt": "4966dc30a35f7884dad1f37d4cf2b14bcf8fdfdbe34f3dbbec1ed2ebd31f9308", + "mappings_256x512/back_left_overlay_mapping.pt": "66796f7e6f321f9bd1739fbd9f2110c6477d814cf9a98f968190a255b2e128e7", + "mappings_256x512/back_right_mapping.pt": "3baa16ca38e7b62f38d04c29746c72c5a3c05ca054d72300cda2c08b2645c4c8", + "mappings_256x512/back_right_overlay_mapping.pt": "63e7dde6cfb22edebe4b45b1ebecaa9ee5679aca1be022e26fabdfc922c84fb8", + "mappings_256x512/front_left_mapping.pt": "f09e38516ed10a81c86939119d776353d795554093c699b10d89741ea78dce7f", + "mappings_256x512/front_left_overlay_mapping.pt": "56db06fe4bf215e06b6f9e8a7b81747e4e24483d770619660b26f555e5402f97", + "mappings_256x512/front_right_mapping.pt": "d9f32bdcbbd98ea6f75aaacc966d6553582d676c06154003dde1559d3f92f1f5", + "mappings_256x512/front_right_overlay_mapping.pt": "138481fd6da1746f21dd90bb42f432a532fa79d73d60515d4558bda2cf56365d", + "mappings_256x512/static_back_mapping.pt": "14213de89eaac720fd83892033c5e42301e682ad37c678dd723619bf10982ca2", + "mappings_256x512/static_front_mapping.pt": "d42e32902fedb9af2bd0f34ae50d18107c3d2264e519beb3b4eb8619d8a0f703", + "mappings_256x512/walk_back_base_layer_ortho_mapping.pt": "9fdc06695caada150b071bc0632fb83a8fed04c6ae1135b25c9519077eddfdca", + "mappings_256x512/walk_back_both_layer_ortho_mapping.pt": "9dfe71bab556365f9e4a33c2eca9d1905806b4278d422adfcf58cb6678b65853", + "mappings_256x512/walk_front_base_layer_ortho_mapping.pt": "c406299c5bab60cc4b8d7afa59f833a41a693ba7b2efa7a819ee2828e4f4008b", + "mappings_256x512/walk_front_both_layer_ortho_mapping.pt": "84ad02344ba9777c18d0ec41004d6fed5211b8a1e5cffd62dcac43ba7e4f3148", + "parser.pt": "ea3bd7d0e0d2eeb73d0bea68233412bf74ddc5ae8ddfb9a3a6957c381f1b2540", + "pipeline.json": "b3ecdcc8400621c21d3d03af249ffa8968b019184b1e82f478581b97ba49a394", + "runtime/LICENSE": "8486a10c4393cee1c25392769ddd3b2d6c242d6ec7928e1414efff7dfb2f07ef", + "runtime/dense_uv_parser/__init__.py": "044e6ec3019c5cf9175ddb7a5df3083630a491d01864837ee6139b01600c2ceb", + "runtime/dense_uv_parser/accessories.py": "5e1787c41640196e94b7a99576d269408f56c80406577cfd0b1afe23ccd8d014", + "runtime/dense_uv_parser/accessory_data.py": "7989807a1771b00df4d80cab9a7f3a0e80a8aa24546c729a9627919ffbab0728", + "runtime/dense_uv_parser/accessory_pipeline.py": "66f8927b310aed4b1dae1eedfb01c1493f8aa7f7bf25e646b662bbe3f40d7ca1", + "runtime/dense_uv_parser/accessory_regression.py": "9fa90be8bd6f3ba50438ec7abc394824a300270fa9dcdf3ad55968a5be73a238", + "runtime/dense_uv_parser/batch_accessories.py": "9800436d24e06e3856abf4e2325546c11ba2b5b97b55099693b3d77d088ae694", + "runtime/dense_uv_parser/cache_semantic_features.py": "9c439ba5a95c5a3f55c8114975872c4ad0eb38087f44f5d2f82e41d07b4a0961", + "runtime/dense_uv_parser/crown_geometry.py": "608a28c8223cb0ffb0882c3ff183b9a44a3f6d2987b660517548ca657cfca14a", + "runtime/dense_uv_parser/foreground.py": "6d488ef64135fd57cc661bcab531748bc242b4034b07b748bf61d34deba61db6", + "runtime/dense_uv_parser/foreground_acceptance.py": "d9c6b9367cc0a4333fe08832bf558a1738a0e533dab257a3be0a181481dfb2ed", + "runtime/dense_uv_parser/foreground_batch.py": "7d7917b61a3dd13a08a5afb181a9c43ad6d8973fb1e485d9503fe79d950730e4", + "runtime/dense_uv_parser/foreground_model.py": "96b1cf0a4944df39d3cb1e3c0cf490508e65306b8b2428fc64fdd6974ffb978a", + "runtime/dense_uv_parser/foreground_provider.py": "1970ad563d3e8aeb7df5e84fc5f4264353d4ceca80f362a13fba75f5cdbeaa5e", + "runtime/dense_uv_parser/head_inpainting.py": "076c8844db555f168599e06a77efa4ae4da48b9d34860a3a2cfdf1da3ccc0638", + "runtime/dense_uv_parser/headwear.py": "0e553177b9214fe6252df3dc9a0ae351c605c7d8127b1de976a04dda4ee5f09d", + "runtime/dense_uv_parser/headwear_data.py": "114f904bd0743206223cf98b71bb027513c2e81c62a1ca0540d5af441384a915", + "runtime/dense_uv_parser/infer.py": "819d6c78b4065b1835a5c39c152b2014adfabf7dfb0180e151b434293d1ae678", + "runtime/dense_uv_parser/losses.py": "87de937f9087284de4a84ce555229c5995ed582d91af99685a2861bc6e8b66f6", + "runtime/dense_uv_parser/material_refine.py": "c236307c541798ef7e232107a9da7d95b95866584c645551477bbce45eca7e18", + "runtime/dense_uv_parser/matting_data.py": "107db91d98bca7a02c7501b648860e75faf7ef32a89919f4cbcd107c479d0429", + "runtime/dense_uv_parser/model.py": "9729d3c97aa4c593e3891609f467563b037bf8dbc9dce597346b5b09bbc5b885", + "runtime/dense_uv_parser/ownership.py": "1f414be584390f0f853423f47af85cf0a5b6a40c1c4503a0e0b76557e70f1ee8", + "runtime/dense_uv_parser/ownership_data.py": "d4ff687cbbba2634aca13f4a9c02db8cf0479b3b2a1bcdc40e9bc501676c6ef1", + "runtime/dense_uv_parser/prepare_foreground_data.py": "13c03bc989b1c89bf2446147b6dd5ef6d0a72bde7e3f903045ef4c7e825c2e91", + "runtime/dense_uv_parser/released_inference.py": "3ca2f766ee6d89788973ef9d116b3b406e2e3556827efaef224d46e71b47cc48", + "runtime/dense_uv_parser/run_local.py": "40c63041f212656f50c0a2e45524e14a2bc09b12b41adb0a0300121e0dd1a039", + "runtime/dense_uv_parser/runtime.py": "c374381063f7f559dc9e2a384982b411c4eb0f5160c0bf0e40f53e4286a15d83", + "runtime/dense_uv_parser/semantic.py": "7d390f827e423d4e97ec9e20e39604c5a4801d97b1796b036016bae55c755249", + "runtime/dense_uv_parser/semantic_backbone.py": "1008138e58f4bc665930d786dc68b6bb44b91bef69bcc71b8d41e252957bc273", + "runtime/dense_uv_parser/semantic_cache.py": "057629ce8aafbf60b379fef487a8a0ba33a2fe7e24a4e11b71ebdb3084489d40", + "runtime/dense_uv_parser/semantic_generalization.py": "c8266c0ffd3a2e07cc1ed42ce62de69d80440c3bbba69b466d1121260d200c0e", + "runtime/dense_uv_parser/semantic_targets.py": "f5e8dd1c7fb53847c5645b8f1cc03fc4594805de6b7d0bf482f223956a5c2008", + "runtime/dense_uv_parser/simple_inpainting.py": "00359b7d9217ccf7b86f2684caee8a464d6142d35b930b464e0e963a74b750ff", + "runtime/dense_uv_parser/skin_dataset.py": "67af52006b15b86c82dc10297ee3e78bc9e634ea20515202c46a6dcb29bb4887", + "runtime/dense_uv_parser/train.py": "76343eb77af5b67d6baefdb80e14cf097da847bb536d47f833f4cc0963ec9233", + "runtime/dense_uv_parser/train_accessories.py": "37de0c34887a4da47f63d7611a5f12a25c45f45e5eea21c6483f5e7b5f4a2a43", + "runtime/dense_uv_parser/train_foreground.py": "6830de93089ef19d5f0e9171bbba19d8ec7dc53d488fd95f0c5346d67883f6f6", + "runtime/dense_uv_parser/train_headphone_presence.py": "b95332c38bce2af876987342b9a3b42fc587722a2531d9a9f3c4f6d0bc6fdc61", + "runtime/dense_uv_parser/train_headwear.py": "61042a2373cecddcc623f42418778b041c481442cc45da098e2c5b27175b8ac1", + "runtime/dense_uv_parser/train_headwear_presence.py": "32249430cf1333e54f1d6915fad32976108ffa720e5ae9f79cd6c571bd2736ac", + "runtime/dense_uv_parser/train_ownership.py": "6d06464752fcf4722d94d6a353c907a8d5cc70bd632ff0253a3ffa1d942ce1ec", + "runtime/dense_uv_parser/utils.py": "34ef3658087014d8a5de36231e4911eb752c4a62fae6601e60725173fed956d3", + "runtime/dense_uv_parser/uv_layout.py": "ddce13bf72bfa01c1e553f1f2760ce5a3ba4b55bae884740a9a3246fe6694e90", + "runtime/dense_uv_parser/uv_topology.py": "8cfe03e229d6f3ea7297165c0f97f5e9b6cf297d7a86da1b6f51c2e6a2f397e2", + "runtime/renderer.py": "e843c16b0d4b2534b6fe24646d15206c24f3aacf820f3fe30b5d42773e199d24", + "tested_environments.json": "fce93e20cf26758a5a7472c1354ba9c2c158d68b3c514b1c19249637033b5bf1" + } +} diff --git a/SKING_DDJ_v101c/runtime/LICENSE b/SKING_DDJ_v101c/runtime/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..0ad25db4bd1d86c452db3f9602ccdbe172438f52 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/LICENSE @@ -0,0 +1,661 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published + by the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/__init__.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a44944eadeac3fc40c3ac8f965a24f30690f9fc --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/__init__.py @@ -0,0 +1,2 @@ +"""Dense render-pixel to Minecraft UV parser.""" + diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/accessories.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessories.py new file mode 100644 index 0000000000000000000000000000000000000000..a667599ab689746e7891ed525b02e801818439eb --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessories.py @@ -0,0 +1,199 @@ +"""Learned head accessory masks, object losses and geometry-safe routing. + +Classes are object identities for one player: none, glasses, hat, outer hair. +RGB colour and projected-cell occupancy never create an accessory label. +""" +import torch +from torch import nn +import torch.nn.functional as F + +ACCESSORY_CLASSES = ('none', 'glasses', 'hat', 'outer_hair') +HEAD_CROP = (0.0625, 0.0, 0.9375, 0.34375) # fixed Steve camera ROI, not a mask + + +def head_bounds(height, width): + x0, y0, x1, y1 = HEAD_CROP + return round(y0*height), round(y1*height), round(x0*width), round(x1*width) + + +def head_crop(images, foreground=None): + y0, y1, x0, x1 = head_bounds(*images.shape[-2:]) + rgb = images[:, :3] + if foreground is None: + foreground = torch.ones_like(rgb[:, :1]) + elif foreground.dim() == 3: + foreground = foreground[:, None] + foreground = foreground.float() + rgb = rgb * foreground + 0.5 * (1-foreground) + crop = torch.cat([rgb, foreground], 1)[:, :, y0:y1, x0:x1] + return F.interpolate(crop.float(), (224, 224), mode='bilinear', align_corners=False) + + +def restore_logits(logits, size): + h, w = size + y0, y1, x0, x1 = head_bounds(h, w) + patch = F.interpolate(logits, (y1-y0, x1-x0), mode='bilinear', align_corners=False) + # Outside the crop the new branch is exactly inactive. + prior = logits.new_tensor([0.]+[-20.]*(logits.shape[1]-1))[None, :, None, None] + return F.pad(patch-prior, (x0, w-x1, y0, h-y1)) + prior + + +def block(cin, cout): + return nn.Sequential(nn.Conv2d(cin, cout, 3, padding=1), nn.GroupNorm(8, cout), + nn.SiLU(), nn.Conv2d(cout, cout, 3, padding=1), + nn.GroupNorm(8, cout), nn.SiLU()) + + +class HeadAccessoryHead(nn.Module): + def __init__(self, semantic_dim=768, predict_hat_components=False): + super().__init__() + self.enc0 = block(4, 24) + self.enc1 = block(24, 48) + self.enc2 = block(48, 96) + self.enc3 = block(96, 96) + self.semantic = nn.Conv2d(semantic_dim, 96, 1) + self.context = nn.TransformerEncoder( + nn.TransformerEncoderLayer(96, 4, 256, 0.05, batch_first=True, norm_first=True), + 2, enable_nested_tensor=False) + self.dec2 = block(192, 96) + self.dec1 = block(144, 48) + self.dec0 = block(72, 24) + self.classifier = nn.Conv2d(24, 4, 1) + self.components = nn.Conv2d(24,6,1) if predict_hat_components else None + if self.components is not None: + nn.init.zeros_(self.components.weight) + nn.init.constant_(self.components.bias,-2.) + with torch.no_grad():self.components.bias[0]=2. + nn.init.constant_(self.classifier.bias, -2.) + with torch.no_grad(): + self.classifier.bias[0] = 2. + + def forward(self, crop, semantic_features, return_components=False): + s0 = self.enc0(crop) + s1 = self.enc1(F.avg_pool2d(s0, 2)) + s2 = self.enc2(F.avg_pool2d(s1, 2)) + z = self.enc3(F.avg_pool2d(s2, 2)) + z = F.avg_pool2d(z, 2) + self.semantic(semantic_features.to(z.dtype)) + z = self.context(z.flatten(2).transpose(1, 2)).transpose(1, 2).reshape_as(z) + z = self.dec2(torch.cat([F.interpolate(z, s2.shape[-2:], mode='bilinear', align_corners=False), s2], 1)) + z = self.dec1(torch.cat([F.interpolate(z, s1.shape[-2:], mode='bilinear', align_corners=False), s1], 1)) + z = self.dec0(torch.cat([F.interpolate(z, s0.shape[-2:], mode='bilinear', align_corners=False), s0], 1)) + logits=self.classifier(z) + return (logits,self.components(z)) if return_components and self.components is not None else logits + + +def accessory_loss(logits, labels, valid=None): + """Per-object Dice + class-balanced CE + boundary affinity, preserving holes.""" + if valid is None: + valid = torch.ones_like(labels, dtype=torch.bool) + target = labels.masked_fill(~valid, -100) + ce = F.cross_entropy(logits.float(), target, weight=logits.new_tensor([1., 3., 2., 2.]).float()) + probability = logits.float().softmax(1) + onehot = F.one_hot(labels.clamp(0, 3), 4).permute(0, 3, 1, 2).float() + mask = valid[:, None] + p, g = probability * mask, onehot * mask + intersection = (p * g).sum((2, 3)) + dice = (1-(2*intersection+1)/(p.sum((2, 3))+g.sum((2, 3))+1))[:, 1:].mean() + # Neighbour agreement targets also include object/background boundaries; + # there is no indiscriminate smoothness term that fills lens/frame holes. + affinity = logits.sum() * 0 + for axis in (2, 3): + delta = torch.diff(p[:, 1:], dim=axis) + truth_delta = torch.diff(g[:, 1:], dim=axis) + affinity = affinity + F.smooth_l1_loss(delta, truth_delta) + return ce + dice + .2 * affinity + + +def connected_object_support(probability, identity, valid, seed_threshold=.90, grow_threshold=0.0): + """Keep a learned object region around confident seeds, across RGB changes. + + Growth is restricted to the SAME predicted object class and foreground. + A low-confidence region without a seed cannot manufacture an accessory. + """ + seeds = (probability >= seed_threshold) & (identity > 0) & valid + grown = seeds.clone() + for cls in (1,2,3): + candidate = (identity == cls) & (probability >= grow_threshold) & valid + current = seeds & (identity == cls) + if not current.any(): + continue + while True: + neighbours = current.clone() + neighbours[:,1:] |= current[:,:-1] + neighbours[:,:-1] |= current[:,1:] + neighbours[:,:,1:] |= current[:,:,:-1] + neighbours[:,:,:-1] |= current[:,:,1:] + updated = current | (candidate & neighbours) + if torch.equal(updated,current): + break + current = updated + grown |= current + return grown, seeds + + +def apply_accessory_routing(routing, outputs, foreground, renderer, views, threshold=.80): + """Promote learned object evidence even when the old route selected inner. + + This is an explicit object-evidence bypass of local texel voting. Actual + foreground and valid head geometry are required; later silhouette vetoes + and colour-background safeguards still apply. + """ + from SkingToolkit.dense_uv_parser.utils import build_static_surface_routing + threshold = float(outputs.get('accessory_route_threshold', threshold)) + logits = outputs.get('accessory_logits') + support = torch.zeros_like(foreground) + if logits is None: + return support + probability, identity = logits.float().softmax(1).max(1) + valid = torch.zeros_like(foreground) + for vi,view in enumerate(views): + static = build_static_surface_routing(renderer,view,foreground.device) + valid[vi::len(views)] = foreground[vi::len(views)] & static['masks'][1] & (static['part'][1] == 0) + support, seeds = connected_object_support(probability,identity,valid,threshold) + for vi, view in enumerate(views): + sl = slice(vi, foreground.shape[0], len(views)) + static = build_static_surface_routing(renderer, view, foreground.device) + accepted = support[sl] + for name, value in [('layer', 1), ('foreground', True), ('surface', 1), + ('secondary', False), ('secondary_routed', False), + ('semantic_fallback', False), ('consensus_outer_gate_rejected', False)]: + if name in routing: + routing[name][sl] = torch.where(accepted, torch.full_like(routing[name][sl], value), routing[name][sl]) + for name in ('flat_uv', 'part', 'face', 'texel_center_score'): + if name in routing: + routing[name][sl] = torch.where(accepted, static[name][1], routing[name][sl]) + routing['confidence'][sl] = torch.where(accepted, probability[sl], routing['confidence'][sl]) + margin = (2*probability[sl]-1).clamp_min(0) + for name, value in [('confidence_margin', margin), ('confidence_margin_ratio', margin/probability[sl].clamp_min(1e-6))]: + if name in routing: + routing[name][sl] = torch.where(accepted, value, routing[name][sl]) + # Learned inner crown/band evidence can also correct an old outer route. + # This is restricted to explicitly trained components; no RGB rule is used. + inner_support=torch.zeros_like(foreground) + component_logits=outputs.get('hat_component_logits') + if component_logits is not None: + cp,ci=component_logits.float().softmax(1).max(1) + candidate=((ci==1)|(ci==2))&(cp>=.95)&(identity==0)&foreground + for vi,view in enumerate(views): + sl=slice(vi,foreground.shape[0],len(views)) + static=build_static_surface_routing(renderer,view,foreground.device) + accepted=candidate[sl]&static['masks'][0]&(static['part'][0]==0) + inner_support[sl]=accepted + for name,value in [('layer',0),('foreground',True),('surface',0),('secondary',False),('secondary_routed',False),('semantic_fallback',False)]: + if name in routing:routing[name][sl]=torch.where(accepted,torch.full_like(routing[name][sl],value),routing[name][sl]) + for name in ('flat_uv','part','face','texel_center_score'): + if name in routing:routing[name][sl]=torch.where(accepted,static[name][0],routing[name][sl]) + routing['confidence'][sl]=torch.where(accepted,cp[sl],routing['confidence'][sl]) + cm=(2*cp[sl]-1).clamp_min(0) + for name,value in [('confidence_margin',cm),('confidence_margin_ratio',cm/cp[sl].clamp_min(1e-6))]: + if name in routing:routing[name][sl]=torch.where(accepted,value,routing[name][sl]) + routing['hat_inner_supported']=inner_support + routing['accessory_identity'] = identity + routing['accessory_probability'] = probability + routing['accessory_supported'] = support + routing['accessory_seed'] = seeds + from SkingToolkit.dense_uv_parser.ownership import apply_head_ownership + apply_head_ownership(routing, outputs, foreground, renderer, views) + from SkingToolkit.dense_uv_parser.headwear import apply_headwear_routing + apply_headwear_routing(routing, outputs, foreground, renderer, views) + return routing['accessory_supported'] diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_data.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_data.py new file mode 100644 index 0000000000000000000000000000000000000000..b082698182bd9f22de48fc9b2da55d3504ae67bb --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_data.py @@ -0,0 +1,148 @@ +"""Procedural object-labelled skins. Labels describe authored objects, not RGB rules.""" +import numpy as np +import torch +from torch.utils.data import Dataset +from SkingToolkit.dense_uv_parser.skin_dataset import load_skin + + +def make_accessory_skin(original, seed, return_components=False): + rng = np.random.default_rng(seed) + uv = original.permute(1, 2, 0).numpy().copy() + objects = np.zeros((64, 64), dtype=np.int64) + components = np.zeros((64,64),dtype=np.int64) # none, inner crown/band, outer crown/band, brim + # All six head faces are authored together. Body pixels remain the source skin. + uv[:16, :32] = 0 + uv[:16, 32:] = 0 + faces = [(8, 8), (24, 8), (0, 8), (16, 8), (8, 0), (16, 0)] + skin = np.array(rng.choice([[.87,.63,.43],[.45,.28,.19],[.65,.44,.30], + [.95,.78,.62],[.32,.20,.13],[.74,.62,.49]])) + if rng.random() < .12: + skin = rng.uniform(.2, .9, 3) + hair = rng.uniform(.03, .30, 3) if rng.random() < .7 else rng.uniform(.1, .85, 3) + for fi, (x, y) in enumerate(faces): + tex = np.clip(skin[None,None] + rng.normal(0,.025,(8,8,1)), 0, 1) + if fi in (1,4): tex[:] = hair + rng.normal(0,.025,(8,8,1)) + if fi in (0,2,3): tex[:rng.integers(1,3)] = hair + uv[y:y+8, x:x+8, :3] = np.clip(tex,0,1) + uv[y:y+8, x:x+8, 3] = 1 + # Eyes, eyebrows and mouth are explicit negative facial examples. + for x in (9, 13): + uv[11:12, x:x+2, :3] = rng.uniform(.01,.18,3) + uv[14,11:13,:3] = rng.uniform(.05,.5,3) + if rng.random() < .4: + uv[10,9:11,:3] = hair + uv[10,13:15,:3] = hair + + def put(face, mask, color, identity, layer=1, component=0): + x, y = faces[face]; x += 32*layer + patch = uv[y:y+8, x:x+8] + patch[mask, :3] = color[mask] if np.asarray(color).ndim == 3 else color + patch[mask, 3] = 1 + objects[y:y+8, x:x+8][mask] = identity + components[y:y+8, x:x+8][mask] = component + + yy, xx = np.mgrid[:8,:8] + negative = rng.random() < .18 + if not negative and rng.random() < .78: + top = int(rng.integers(2,4)); height = int(rng.integers(3,5)) + frame_color = rng.uniform(.025,.23,3) if rng.random()<.70 else rng.uniform(.01,.85,3) + lens_color = rng.uniform(.05,.95,3) + frame = np.zeros((8,8),bool); lenses = np.zeros_like(frame) + # Distinct connected eyewear styles; class ID is independent of palette. + visor = rng.random()<.20 + boxes = [(0,3),(4,7)] if visor else [(0,2),(5,7)] + hollow = rng.random()<.30 + rimless_top = rng.random()<.45 and not hollow + for box_index,(left, right) in enumerate(boxes): + rectangle = (xx>=left)&(xx<=right)&(yy>=top)&(yy.45)) + tex=np.clip(np.broadcast_to(hair,(8,8,3))+rng.normal(0,.025,(8,8,1)),0,1) + put(face,mask,tex,3) + if return_components: + return torch.from_numpy(uv).permute(2,0,1).float(),torch.from_numpy(objects),torch.from_numpy(components) + return torch.from_numpy(uv).permute(2,0,1).float(), torch.from_numpy(objects) + + +class AccessoryDataset(Dataset): + def __init__(self, paths, count, seed): + self.paths, self.count, self.seed, self.epoch = list(paths), count, seed, 0 + + def __len__(self): return self.count + + def __getitem__(self, index): + seed = self.seed + index*104729 + self.epoch*self.count*104729 + source = load_skin(self.paths[(index + self.epoch*self.count) % len(self.paths)]) + uv, objects, components = make_accessory_skin(source, seed, return_components=True) + return {'uv':uv,'objects':objects,'components':components,'seed':seed} + + +def render_accessories(uv, objects, renderer, views, components=None): + from SkingToolkit.dense_uv_parser.utils import build_dense_parser_batch + images, labels, foregrounds, parts = [], [], [], [] + for view in views: + rendered, target = build_dense_parser_batch(uv, renderer, view) + xy = (target['uv']*63).round().long().clamp(0,63) + flat = xy[:,1]*64+xy[:,0] + identity = objects.flatten(1).gather(1,flat.flatten(1)).reshape_as(flat) + identity = identity.masked_fill((target['layer'] != 1) | (target['foreground'][:,0] < .5), 0) + images.append(rendered);labels.append(identity);foregrounds.append(target['foreground'][:,0].bool()) + if components is not None: + part=components.flatten(1).gather(1,flat.flatten(1)).reshape_as(flat) + parts.append(part.masked_fill(target['foreground'][:,0]<.5,0)) + b = uv.shape[0] + result=(torch.stack(images,1).flatten(0,1),torch.stack(labels,1).flatten(0,1),torch.stack(foregrounds,1).flatten(0,1)) + return result+(torch.stack(parts,1).flatten(0,1),) if components is not None else result diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_pipeline.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..0e42196ad87b7c50820a86a056bbeb722e17c04b --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_pipeline.py @@ -0,0 +1,89 @@ +"""The same foreground, routing, completion and renderer path for v101 eval/infer.""" +import json +from pathlib import Path +import torch +from SkingToolkit.dense_uv_parser.foreground import build_parser_input,learned_foreground_support +from SkingToolkit.dense_uv_parser.utils import estimate_top_left_flood_foreground, splat_parser_predictions_to_uv_conditioning +from SkingToolkit.dense_uv_parser.infer import simple_inpaint_uv + + +def load_pipeline(path=None): + return json.loads(Path(path or Path(__file__).with_name('v101_pipeline.json')).read_text()) + + +@torch.no_grad() +def run_pipeline(model, renderer, images, config=None, complete=False, outputs=None, foreground_probability=None): + config = config or load_pipeline() + views = config['routing']['views'] + color_sources=None + if foreground_probability is None: + fg = estimate_top_left_flood_foreground(images, config['foreground_flood_tolerance']) + else: + if foreground_probability.shape != images.shape[:1]+images.shape[-2:]: + raise ValueError('Foreground probability must have shape (N,H,W)') + probability=foreground_probability.to(device=images.device,dtype=torch.float32) + if not torch.isfinite(probability).all() or probability.min()<0 or probability.max()>1: + raise ValueError('Foreground probability must be finite in [0,1]') + fg,color_sources=learned_foreground_support(probability,config.get('foreground_probability_threshold',.5),config.get('foreground_source_threshold',.98),config.get('foreground_source_inset',1)) + actual = build_parser_input(images, fg, background_mode=config['foreground_background']) + if outputs is None: + outputs = model(actual, view_ids=torch.arange(images.shape[0], device=images.device) % len(views), semantic_foreground=fg) + outputs = {**outputs, 'accessory_route_threshold':config['accessory_route_threshold']} + outputs['headwear_presence_threshold']=config.get('headwear_presence_threshold',.95) + outputs['headphone_presence_threshold']=config.get('headphone_presence_threshold',.9) + outputs['headphone_presence_consensus']=config.get('headphone_presence_consensus',False) + mode=config.get('headphone_routing_mode','unrestricted') + def splat(current): + return splat_parser_predictions_to_uv_conditioning( + images,current,renderer=renderer,observed_foreground=fg,observed_color_support=color_sources,return_details=True,**config['routing']) + if mode=='existing_uv' and 'head_ownership_logits' in outputs: + # The new headphone head may recover colour ownership inside geometry + # already accepted by v101. It cannot manufacture new outer occupancy. + initial,_=splat({**outputs,'headphone_routing_mode':'disabled'}) + outputs={**outputs,'headphone_routing_mode':'existing_uv','headphone_uv_support':(initial[:,9]>.5).flatten(1)} + else:outputs={**outputs,'headphone_routing_mode':mode} + cond,details=splat(outputs) + geometry_mode = config.get('crown_top_geometry_mode', 'legacy_consensus') + if geometry_mode == 'rendered_semantics': + from SkingToolkit.dense_uv_parser.crown_geometry import reconcile_crown_geometry + cond = reconcile_crown_geometry(cond, details, renderer, views) + elif geometry_mode == 'legacy_consensus': + from SkingToolkit.dense_uv_parser.headwear import reconcile_crown_top + cond = reconcile_crown_top(cond, details, renderer, views) + else: + raise ValueError('Unknown crown top geometry mode: ' + str(geometry_mode)) + result = {'foreground_color_sources':color_sources,'foreground_probability':foreground_probability,'conditioning':cond,'details':details,'outputs':outputs,'foreground':fg} + if complete: + result['uv'] = torch.stack([simple_inpaint_uv(c[None].cpu())[0] for c in cond]).to(images.device) + if config.get('head_semantic_inpaint',False): + from SkingToolkit.dense_uv_parser.head_inpainting import repair_hidden_head_material + result['uv']=repair_hidden_head_material(result['uv'],cond,details,views) + from SkingToolkit.dense_uv_parser.headwear import isolate_headwear_material + result['uv'],headwear_uv=isolate_headwear_material(result['uv'],cond,details,views) + if config.get('head_material_refine_steps',0): + from SkingToolkit.dense_uv_parser.material_refine import refine_head_material + if config.get('head_material_refine_mode')!='isolated': + from SkingToolkit.dense_uv_parser.material_refine import refine_head_material_legacy + result['uv']=refine_head_material_legacy(result['uv'],details['rendered'],details['color_source_support'],renderer,views,config['head_material_refine_steps']) + else: + result['uv']=refine_head_material(result['uv'],details['rendered'],details['color_source_support'],renderer,views,config['head_material_refine_steps'],ownership=details['outputs'].get('head_ownership_logits'),inner_exclusion=details['routing'].get('head_color_boundary_excluded'),inner_texel_support=cond[:,4]>.5,headwear_layer=details['routing'].get('headwear_layer'),headwear_uv=headwear_uv,headwear_family=details['routing'].get('headwear_family')) + result['render'] = torch.stack([renderer.forward_view(result['uv'], view) for view in views],1).flatten(0,1) + return result + + +def cached_real_foreground(images, source, config): + """Load verified learned foreground for real development evaluation.""" + import hashlib + import numpy as np + from PIL import Image + import torch.nn.functional as F + cache=config.get('foreground_cache') + if cache is None:return None + source=Path(source).resolve();entry=cache.get(str(source)) + if entry is None:raise ValueError('Missing learned foreground for '+str(source)) + data=Path(entry['probability']).read_bytes() + if hashlib.sha256(source.read_bytes()).hexdigest()!=entry['input_sha256'] or hashlib.sha256(data).hexdigest()!=entry['probability_sha256']: + raise ValueError('Learned foreground cache changed') + import io + p=torch.from_numpy(np.array(Image.open(io.BytesIO(data))).copy()).float().to(images.device)/255. + return F.interpolate(torch.stack(p.chunk(len(config['routing']['views']),1))[:,None],images.shape[-2:],mode='nearest-exact')[:,0] diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_regression.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_regression.py new file mode 100644 index 0000000000000000000000000000000000000000..1b3311a0356c182940a26e287c78586f0627f818 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/accessory_regression.py @@ -0,0 +1,73 @@ +"""Visually traced development regression, kept separate from training/test data.""" +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +import numpy as np +from PIL import Image +import torch +from PIL import ImageDraw,ImageFilter +import torch.nn.functional as F +from torchvision.utils import save_image +from SkingToolkit.dense_uv_parser.infer import load_view_images +from SkingToolkit.dense_uv_parser.accessory_pipeline import run_pipeline,cached_real_foreground +from SkingToolkit.dense_uv_parser.uv_layout import tensor_to_rgba_image + + +@torch.no_grad() +def evaluate_real_regression(model,renderer,config,output_dir,step): + root=Path(__file__).parent + metadata=json.loads((root/'v101_real_regression.json').read_text()) + image_path=Path(metadata['combined']) + if hashlib.sha256(image_path.read_bytes()).hexdigest()!=metadata['source_sha256']: + raise ValueError('Regression input changed; review the annotation before using it') + images=load_view_images(SimpleNamespace(combined=image_path),config['routing']['views'],renderer).cuda() + probability=cached_real_foreground(images,image_path,config) + result=run_pipeline(model,renderer,images,config,complete=True,foreground_probability=probability) + baseline=run_pipeline(model,renderer,images,config,outputs={k:v for k,v in result['outputs'].items() if k not in ('accessory_logits','hat_component_logits')},foreground_probability=probability) + annotation=torch.from_numpy(np.array(Image.open(root/metadata['mask']))>0).cuda() + eroded=(-F.max_pool2d(-annotation[None,None].float(),5,1,2))[0,0]>.5 + expanded=F.max_pool2d(annotation[None,None].float(),5,1,2)[0,0]>.5 + uncertain=expanded&~eroded + pred=result['outputs']['accessory_logits'][0].softmax(0) + glasses=(pred.argmax(0)==1)&result['details']['routing']['accessory_supported'][0] + valid=result['foreground'][0]&~uncertain + tp=int((glasses&annotation&valid).sum());fp=int((glasses&~annotation&valid).sum());fn=int((~glasses&annotation&valid).sum()) + r=result['details']['routing'];br=baseline['details']['routing'] + final_outer=r['foreground'][0]&(r['layer'][0]==1) + base_outer=br['foreground'][0]&(br['layer'][0]==1) + rgb=images[0,:3];yy=torch.arange(rgb.shape[1],device=rgb.device)[:,None];xx=torch.arange(rgb.shape[2],device=rgb.device)[None,:] + green=(rgb[1]>rgb[0]+.12)&(rgb[1]>rgb[2]+.14)&(rgb[1]>.25)&(yy<145)&(xx>90) + metrics={'scope':metadata['role'],'glasses_mask_precision':tp/max(1,tp+fp),'glasses_mask_recall':tp/max(1,tp+fn), + 'whole_glasses_outer_recall':float(final_outer[eroded].float().mean()), + 'baseline_whole_glasses_outer_recall':float(base_outer[eroded].float().mean())} + for name,mask in [('left',green&(xx<160)),('right',green&(xx>=160))]: + metrics[name+'_lens_outer_recall']=float(final_outer[mask].float().mean()) + metrics['baseline_'+name+'_lens_outer_recall']=float(base_outer[mask].float().mean()) + save_image(torch.cat(list(result['render'][:,:3]),dim=2).cpu(),output_dir/f'real_{step}_render.png') + save_image(pred[1:2].cpu(),output_dir/f'real_{step}_glasses_probability.png') + tensor_to_rgba_image(result['uv'][0]).save(output_dir/f'real_{step}_uv.png') + if config.get("hat_development"): + metrics.update(evaluate_hat_regression(model,renderer,config,output_dir,step)) + return metrics + + +@torch.no_grad() +def evaluate_hat_regression(model,renderer,config,output_dir,step): + spec=config['hat_development'];source=Path(spec['input']) + if hashlib.sha256(source.read_bytes()).hexdigest()!=spec['input_sha256']:raise ValueError('Hat input changed') + images=load_view_images(SimpleNamespace(combined=source),config['routing']['views'],renderer).cuda() + probability=cached_real_foreground(images,source,config) + result=run_pipeline(model,renderer,images,config,complete=True,foreground_probability=probability) + mask=torch.from_numpy(np.array(Image.open(spec['brim_mask']).filter(ImageFilter.MinFilter(3)))>0).cuda() + routing=result['details']['routing'];outer=routing['foreground']&(routing['layer']==1) + # Explicit colour-band region, used only as a development render diagnostic. + band=Image.new('L',(256,512));ImageDraw.Draw(band).polygon([(48,56),(107,80),(207,65),(207,79),(107,94),(48,69)],fill=255) + band=torch.from_numpy(np.array(band.filter(ImageFilter.MaxFilter(9)))>0).cuda() + def chroma(rgb):return (rgb[:,0]-rgb[:,1:3].amax(1)-.15).clamp_min(0)[:,band].sum(1) + reference=chroma(images[:,:3]);observed=chroma(result['render'][:,:3]) + metrics={'hat_brim_outer_recall':[float(v[mask].float().mean()) for v in outer], + 'hat_band_chroma_retention':[float(x) for x in (observed/reference.clamp_min(1e-5))]} + save_image(torch.cat(list(result['render'][:,:3]),2).cpu(),output_dir/f'hat_{step}_render.png') + tensor_to_rgba_image(result['uv'][0]).save(output_dir/f'hat_{step}_uv.png') + return metrics diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/batch_accessories.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/batch_accessories.py new file mode 100644 index 0000000000000000000000000000000000000000..e777aabfce833ddea3dfe35e1c9842c0563ce84a --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/batch_accessories.py @@ -0,0 +1,147 @@ +"""Fingerprint-addressed batch inference with isolated, atomic sample outputs.""" +import argparse +import hashlib +import json +import io +import os +import shutil +import tempfile +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace +import torch +import numpy as np +from PIL import Image +import torch.nn.functional as F +from torchvision.utils import save_image +from SkingToolkit.dense_uv_parser.infer import load_parser, load_view_images, save_parser_uv +from SkingToolkit.dense_uv_parser.accessory_pipeline import load_pipeline, run_pipeline +from SkingToolkit.dense_uv_parser.uv_layout import tensor_to_rgba_image +from SkingToolkit.renderer import DifferentiableRenderer + + +def sha(path): + with Path(path).open('rb') as f:return hashlib.file_digest(f,'sha256').hexdigest() + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('--checkpoint',required=True,type=Path) + parser.add_argument('--output-dir',required=True,type=Path) + parser.add_argument('--inputs',nargs='+',type=Path) + parser.add_argument('--pipeline',type=Path) + parser.add_argument('--save-tensors',action='store_true') + parser.add_argument('--foreground-probabilities',type=Path,nargs='+',help='One combined grayscale probability image per input; bypass flood only when explicitly supplied') + opt=parser.parse_args();torch.set_num_threads(4) + config=load_pipeline(opt.pipeline) + cases=opt.inputs or [Path(x) for x in json.loads(Path(__file__).with_name('v101_cases.json').read_text())] + if opt.foreground_probabilities is not None and len(opt.foreground_probabilities)!=len(cases): + parser.error('Provide exactly one foreground probability image for each input') + if not all(x.is_file() for x in cases):raise FileNotFoundError([str(x) for x in cases if not x.is_file()]) + # A training job may atomically replace best.pt while inference is starting. + # Hash and load the same byte snapshot so output provenance stays accurate. + checkpoint_bytes=opt.checkpoint.read_bytes() + checkpoint_sha256=hashlib.sha256(checkpoint_bytes).hexdigest() + model,args=load_parser(io.BytesIO(checkpoint_bytes),torch.device('cuda')) + del checkpoint_bytes + if opt.pipeline is None and args.get('_v101_pipeline'): + config=args['_v101_pipeline'] + if opt.foreground_probabilities is None and config.get('foreground_release'): + # New checkpoints must use their paired trained foreground model even + # when called through the lower-level batch entry point. + root=Path(__file__).resolve().parents[1];release=config['foreground_release'] + delta=root/release['checkpoint'] + if not release['passed'] or sha(delta)!=release['checkpoint_sha256']:raise ValueError('Foreground release changed') + with tempfile.TemporaryDirectory(prefix='foreground-auto-') as temporary: + index=Path(temporary)/'predictions.json' + subprocess.run([sys.executable,str(root/'dense_uv_parser/foreground_provider.py'),'--checkpoint',str(delta),'--model-dir',release['model_dir'],'--output-dir',str(root/'dense_uv_parser/cache/released_foreground'),'--result-manifest',str(index),'--inputs',*[str(x.resolve()) for x in cases]],cwd=root,check=True) + opt.foreground_probabilities=[Path(x) for x in json.loads(index.read_text())['probabilities']] + if len(opt.foreground_probabilities)!=len(cases):raise ValueError('Foreground output count does not match inputs') + renderer=DifferentiableRenderer(args['mappings_dir']).cuda() + opt.output_dir.mkdir(parents=True,exist_ok=True) + sources={p.name:sha(p) for p in Path(__file__).parent.glob('*.py')} + common={'checkpoint':str(opt.checkpoint.resolve()),'checkpoint_sha256':checkpoint_sha256,'pipeline':config,'source_sha256':sources,'save_tensors':opt.save_tensors} + for case_index,case in enumerate(cases): + case_bytes=case.read_bytes() + manifest={**common,'input':str(case.resolve()),'input_sha256':hashlib.sha256(case_bytes).hexdigest()} + probability_bytes=None + if opt.foreground_probabilities is not None: + probability_path=opt.foreground_probabilities[case_index] + probability_bytes=probability_path.read_bytes() + sidecar=probability_path.parent/'manifest.json' + if sidecar.is_file(): + provider=json.loads(sidecar.read_text()) + if provider.get('input_sha256')!=manifest['input_sha256']:raise ValueError('Foreground prediction belongs to different input bytes') + manifest['foreground_provider']={k:provider.get(k) for k in ('model_sha256','adaptation_sha256','adaptation_version','provider_sha256')} + manifest['foreground_probability']={'path':str(probability_path.resolve()),'sha256':hashlib.sha256(probability_bytes).hexdigest(),'threshold':config.get('foreground_probability_threshold',.5),'source_threshold':config.get('foreground_source_threshold',.98),'source_inset':config.get('foreground_source_inset',1)} + fingerprint=hashlib.sha256(json.dumps(manifest,sort_keys=True).encode()).hexdigest() + # Include both path/input/checkpoint hashes: same basenames cannot collide. + target=opt.output_dir/(case.stem+'_'+fingerprint[:12]) + if (target/'manifest.json').is_file(): + old=json.loads((target/'manifest.json').read_text()) + if old.get('fingerprint')==fingerprint and old.get('complete'): + print('cached='+str(target),flush=True);continue + if target.exists():raise FileExistsError('Incomplete or conflicting output: '+str(target)) + scratch=Path(tempfile.mkdtemp(prefix='.inference-',dir=opt.output_dir)) + try: + images=load_view_images(SimpleNamespace(combined=io.BytesIO(case_bytes)),config['routing']['views'],renderer).cuda() + probability=None + if probability_bytes is not None: + raw=Image.open(io.BytesIO(probability_bytes)) + if raw.mode not in ('L','F','I;16'):raise ValueError('Foreground probability must be a grayscale image') + with Image.open(io.BytesIO(case_bytes)) as source_image: + if raw.size!=source_image.size:raise ValueError('Foreground probability must match original combined image dimensions') + p=torch.from_numpy(np.array(raw).astype(np.float32)) + p=p/(65535. if raw.mode=='I;16' else 255.) if raw.mode!='F' else p + if p.shape[1]%len(config['routing']['views']):raise ValueError('Combined width must divide evenly into views') + probability=F.interpolate(torch.stack(p.chunk(len(config['routing']['views']),dim=1))[:,None],images.shape[-2:],mode='nearest-exact')[:,0].cuda() + result=run_pipeline(model,renderer,images,config,complete=True,foreground_probability=probability) + if probability is not None: + save_image(probability[:,None].cpu(),scratch/'foreground_probability.png',nrow=len(config['routing']['views'])) + save_image(result['foreground'][:,None].float().cpu(),scratch/'foreground_mask.png',nrow=len(config['routing']['views'])) + save_image(result['foreground_color_sources'][:,None].float().cpu(),scratch/'foreground_color_sources.png',nrow=len(config['routing']['views'])) + save_parser_uv(result['conditioning'],scratch/'parser_only_uv.png') + for name in ('parser_pred_uv_simple_inpainting.png','pred_uv.png'): + tensor_to_rgba_image(result['uv'][0]).save(scratch/name) + save_image(torch.cat(list(result['render'][:,:3]),dim=2).cpu(),scratch/'simple_inpaint_render.png') + routing=result['details']['routing'] + for layer,name in [(0,'inner'),(1,'outer')]: + mask=routing['foreground']&(routing['layer']==layer) + cutout=torch.where(mask[:,None],images[:,:3],torch.full_like(images[:,:3],.5)) + save_image(cutout.cpu(),scratch/f'parser_debug_{name}.png',nrow=len(config['routing']['views'])) + logits=result['outputs'].get('accessory_logits') + if logits is not None: + palette=images.new_tensor([[.5,.5,.5],[.1,.9,.5],[.95,.6,.1],[.7,.3,.9]]) + save_image(palette[logits.argmax(1)].permute(0,3,1,2).cpu(),scratch/'accessory_mask.png',nrow=len(config['routing']['views'])) + save_image((1-logits.softmax(1)[:,:1]).cpu(),scratch/'accessory_probability.png',nrow=len(config['routing']['views'])) + ownership=result['outputs'].get('head_ownership_logits') + if ownership is not None: + palette=images.new_tensor([[.5,.5,.5],[1.,.65,.4],[.4,.3,.8],[.1,.9,.5],[1.,.2,.3]]) + save_image(palette[ownership.argmax(1)].permute(0,3,1,2).cpu(),scratch/'head_ownership.png',nrow=len(config['routing']['views'])) + headwear=result['outputs'].get('headwear_logits') + if headwear is not None: + palette=images.new_tensor([[.5,.5,.5],[.1,.8,.8],[.1,.3,1],[.8,.2,.8],[1,.3,.2],[1,1,.1],[.95,.65,.05]]) + save_image(palette[headwear.argmax(1)].permute(0,3,1,2).cpu(),scratch/'headwear_components.png',nrow=len(config['routing']['views'])) + components=result['outputs'].get('hat_component_logits') + if components is not None: + palette=images.new_tensor([[.5,.5,.5],[.2,.8,.8],[.2,.4,1.],[.7,.4,.1],[.9,.2,.5],[1.,.85,.1]]) + save_image(palette[components.argmax(1)].permute(0,3,1,2).cpu(),scratch/'hat_components.png',nrow=len(config['routing']['views'])) + if 'crown_geometry' in result['details']: + (scratch/'crown_geometry.json').write_text(json.dumps(result['details']['crown_geometry'], indent=2)) + save_image(result['details']['headwear_removed_top_uv'][:,None].float().cpu(), scratch/'crown_removed_top_uv.png') + # Exact tensors make pixel-stage regression metrics reproducible. + if opt.save_tensors: + torch.save({'images':images.cpu(),'outputs':{k:v.cpu() if torch.is_tensor(v) else v for k,v in result['outputs'].items()}, + 'routing':{k:v.cpu() if torch.is_tensor(v) else v for k,v in routing.items()}, + 'crown_geometry': result['details'].get('crown_geometry'), + 'headwear_removed_top_uv': result['details'].get('headwear_removed_top_uv', torch.zeros(0)).cpu()},scratch/'diagnostics.pt') + manifest.update(fingerprint=fingerprint,complete=True) + (scratch/'manifest.json').write_text(json.dumps(manifest,indent=2)) + os.replace(scratch,target);print('completed='+str(target),flush=True) + except BaseException: + shutil.rmtree(scratch) + raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/cache_semantic_features.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/cache_semantic_features.py new file mode 100644 index 0000000000000000000000000000000000000000..a633cc8b1ad7691872c61ee3a4639c0f3afe4726 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/cache_semantic_features.py @@ -0,0 +1,267 @@ +"""Precompute frozen SigLIP2 features for Dense UV Parser fixed views. + +The direct trainer has no render augmentation, so recomputing the same frozen +vision tower output in every epoch wastes most of the semantic-encoder time. +The optional spatial cache stores full-width FP16 patch features after removing +only deterministic letterbox padding. Trainable projections remain in the +downstream model. +""" + +import argparse +import json +import os +import sys +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +WORKSPACE_ROOT = TOOLKIT_ROOT.parent +if str(WORKSPACE_ROOT) not in sys.path: + sys.path.insert(0, str(WORKSPACE_ROOT)) + +from SkingToolkit.renderer import DifferentiableRenderer # noqa: E402 +from SkingToolkit.dense_uv_parser.semantic_backbone import ( # noqa: E402 + SigLIP2VisionBackbone, +) +from SkingToolkit.dense_uv_parser.semantic_cache import ( # noqa: E402 + SIGLIP_CACHE_VERSION, + SigLIPGlobalCache, +) +from SkingToolkit.dense_uv_parser.runtime import get_device # noqa: E402 +from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset # noqa: E402 +from SkingToolkit.dense_uv_parser.utils import parse_views # noqa: E402 + +try: + from tqdm import tqdm +except ImportError: + tqdm = None + + +def autocast_context(device, precision): + if precision == "no" or device.type == "cpu": + return nullcontext() + dtype = torch.float16 if precision == "fp16" else torch.bfloat16 + return torch.autocast(device_type=device.type, dtype=dtype) + + +def cache_is_reusable( + cache_dir, + dataset, + views, + siglip_model, + require_spatial=False, +): + try: + cache = SigLIPGlobalCache( + cache_dir, + expected_views=views, + expected_model=siglip_model, + expected_data_dir=dataset.data_dir, + require_spatial=require_spatial, + ) + except (FileNotFoundError, ValueError, KeyError, json.JSONDecodeError): + return False + return list(cache.filename_to_index) == [path.name for path in dataset.skin_paths] + + +def build_arg_parser(): + parser = argparse.ArgumentParser( + description="Cache fixed-view SigLIP2 global and optional spatial features." + ) + parser.add_argument("--data_dir", default="../skins") + parser.add_argument("--cache_dir", required=True) + parser.add_argument("--mappings_dir", required=True) + parser.add_argument( + "--views", + default="walk_front_both_layer_ortho,walk_back_both_layer_ortho", + ) + parser.add_argument("--siglip_model", default="google/siglip2-base-patch16-224") + parser.add_argument("--siglip_local_files_only", action="store_true") + parser.add_argument("--max_samples", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=32) + parser.add_argument("--num_workers", type=int, default=16) + parser.add_argument("--prefetch_factor", type=int, default=4) + parser.add_argument("--device", default="auto") + parser.add_argument("--mixed_precision", choices=["no", "fp16", "bf16"], default="bf16") + parser.add_argument( + "--spatial", + action="store_true", + help="Also cache raw cropped patch features as FP16 mmap data.", + ) + parser.add_argument("--force", action="store_true") + return parser + + +@torch.inference_mode() +def main(): + args = build_arg_parser().parse_args() + if args.batch_size < 1: + raise ValueError("--batch_size must be positive.") + views = parse_views(args.views) + if len(views) < 2: + raise ValueError("At least front and back cache views are required.") + dataset = SkinUVDataset(args.data_dir, max_samples=args.max_samples) + cache_dir = Path(args.cache_dir) + if not args.force and cache_is_reusable( + cache_dir, + dataset, + views, + args.siglip_model, + require_spatial=args.spatial, + ): + cache_kind = "global + spatial" if args.spatial else "global" + print(f"Reusing complete SigLIP {cache_kind} cache: {cache_dir}") + return + + device = get_device(args.device) + renderer = DifferentiableRenderer(mappings_dir=args.mappings_dir).to(device).eval() + missing_views = [view for view in views if view not in renderer.views] + if missing_views: + raise ValueError(f"Renderer cache mappings are missing views {missing_views}.") + backbone = SigLIP2VisionBackbone( + model_name=args.siglip_model, + token_channels=128, + local_files_only=args.siglip_local_files_only, + ).to(device).eval() + + loader_kwargs = { + "batch_size": args.batch_size, + "shuffle": False, + "num_workers": args.num_workers, + "pin_memory": device.type == "cuda", + } + if args.num_workers > 0: + loader_kwargs["persistent_workers"] = True + loader_kwargs["prefetch_factor"] = args.prefetch_factor + loader = DataLoader(dataset, **loader_kwargs) + + cache_dir.mkdir(parents=True, exist_ok=True) + temporary_embeddings = cache_dir / ".embeddings.npy.tmp" + temporary_spatial_embeddings = cache_dir / ".spatial_embeddings.npy.tmp" + temporary_metadata = cache_dir / ".metadata.json.tmp" + embeddings = np.lib.format.open_memmap( + temporary_embeddings, + mode="w+", + dtype=np.float16, + shape=(len(dataset), len(views), backbone.raw_feature_dim), + ) + spatial_embeddings = None + spatial_shape = None + iterator = tqdm(loader, desc="cache SigLIP2", leave=False) if tqdm else loader + offset = 0 + for batch in iterator: + uv = batch["uv"].to(device, non_blocking=True) + renders = torch.stack( + [renderer.forward_view(uv, view) for view in views], dim=1 + ) + batch_size, view_count, _, height, width = renders.shape + images = renders[:, :, :3].reshape(batch_size * view_count, 3, height, width) + if images.is_cuda: + images = images.contiguous(memory_format=torch.channels_last) + with autocast_context(device, args.mixed_precision): + encoded = ( + backbone.encode_dense(images) + if args.spatial + else backbone.encode_global(images) + ) + raw_global = encoded["raw_global"].reshape( + batch_size, view_count, backbone.raw_feature_dim + ) + embeddings[offset : offset + batch_size] = ( + raw_global.detach().to(device="cpu", dtype=torch.float16).numpy() + ) + if args.spatial: + raw_spatial = encoded["raw_spatial"] + current_spatial_shape = tuple(raw_spatial.shape[-2:]) + if spatial_embeddings is None: + spatial_shape = current_spatial_shape + spatial_embeddings = np.lib.format.open_memmap( + temporary_spatial_embeddings, + mode="w+", + dtype=np.float16, + shape=( + len(dataset), + len(views), + backbone.raw_spatial_feature_dim, + *spatial_shape, + ), + ) + spatial_size_gib = ( + np.prod(spatial_embeddings.shape) + * np.dtype(np.float16).itemsize + / (1024**3) + ) + print( + "SigLIP2 spatial cache: " + f"shape={spatial_embeddings.shape}, " + f"estimated_size={spatial_size_gib:.1f} GiB." + ) + elif current_spatial_shape != spatial_shape: + raise ValueError( + "SigLIP spatial shape changed within a fixed-view cache: " + f"expected {spatial_shape}, got {current_spatial_shape}." + ) + raw_spatial = raw_spatial.reshape( + batch_size, + view_count, + backbone.raw_spatial_feature_dim, + *spatial_shape, + ) + spatial_embeddings[offset : offset + batch_size] = ( + raw_spatial.detach().to(device="cpu", dtype=torch.float16).numpy() + ) + offset += batch_size + embeddings.flush() + del embeddings + if spatial_embeddings is not None: + spatial_embeddings.flush() + del spatial_embeddings + if offset != len(dataset): + raise RuntimeError(f"Cached {offset} samples, expected {len(dataset)}.") + + metadata = { + "version": SIGLIP_CACHE_VERSION, + "data_dir": str(Path(args.data_dir).resolve()), + "filenames": [path.name for path in dataset.skin_paths], + "views": views, + "siglip_model": args.siglip_model, + "feature_dim": backbone.raw_feature_dim, + "dtype": "float16", + } + if args.spatial: + if spatial_shape is None: + raise RuntimeError("Spatial SigLIP cache was requested but no features were written.") + metadata.update( + { + "spatial_feature_dim": backbone.raw_spatial_feature_dim, + "spatial_height": spatial_shape[0], + "spatial_width": spatial_shape[1], + "spatial_dtype": "float16", + "spatial_layout": "N,V,C,H,W", + "spatial_letterbox_padding_removed": True, + } + ) + temporary_metadata.write_text( + json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8" + ) + os.replace(temporary_embeddings, cache_dir / "embeddings.npy") + if args.spatial: + os.replace( + temporary_spatial_embeddings, + cache_dir / "spatial_embeddings.npy", + ) + os.replace(temporary_metadata, cache_dir / "metadata.json") + cache_kind = "global + spatial" if args.spatial else "global" + print( + f"Cached {cache_kind} features for {len(dataset)} skins x " + f"{len(views)} views to {cache_dir} " + f"({backbone.raw_feature_dim} features/view)." + ) + + +if __name__ == "__main__": + main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/crown_geometry.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/crown_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..c088df6347a4f63ce5baeac64914413d2db99047 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/crown_geometry.py @@ -0,0 +1,129 @@ +"""Resolve crown top occupancy through the renderer's actual visible surfaces. + +Pixel semantics alone cannot distinguish a primary top hit from a crown wall +behind it. Compare binary UV hypotheses through the same occlusion/fallback +renderer used for the final image. RGB and source filenames are not evidence. +""" +import copy + +import torch + +from SkingToolkit.dense_uv_parser.accessories import head_bounds +from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology + + +def crop_renderer(renderer, views, height, width): + """Share immutable buffers, slicing only image axes; never mutate renderer.""" + y0, y1, x0, x1 = head_bounds(height, width) + cropped = copy.copy(renderer) + cropped._buffers = renderer._buffers.copy() + for name, value in cropped._buffers.items(): + if not any(name.startswith(view + '_') for view in views): + continue + for axis in range(value.ndim - 1): + if value.shape[axis:axis + 2] == (height, width): + slices = [slice(None)] * value.ndim + slices[axis] = slice(y0, y1) + slices[axis + 1] = slice(x0, x1) + cropped._buffers[name] = value[tuple(slices)] + break + return cropped, (slice(y0, y1), slice(x0, x1)) + + +@torch.no_grad() +def prune_crown_top(uv, target, foreground, renderer, views, crown_uv=None): + """Delete a top cell only when its removal improves rendered semantics. + + Only existing outer head-top alpha can change. Hidden or ambiguous cells + stay unchanged: no sparsity penalty, fixed open-top template, or minimum + crown thickness. All other geometry and material stay frozen. The binary + search re-evaluates visibility after each deletion, including secondary + crown walls revealed behind the top plane. Each player is independent. + """ + if target.shape[0] != uv.shape[0] * len(views) or foreground.shape != target.shape: + raise ValueError('Crown geometry requires complete aligned view groups') + topology = build_simple_uv_topology() + outer = (topology.valid & (topology.part == 0) & (topology.layer == 1)).to(uv.device) + top = outer & (topology.face.to(uv.device) == 4) + cropped, region = crop_renderer(renderer, views, *target.shape[-2:]) + target = target[..., region[0], region[1]].float() + foreground = foreground[..., region[0], region[1]].float() + result = uv.clone() + records = [] + for group in range(uv.shape[0]): + original = uv[group:group + 1] + semantic = torch.zeros_like(original) + semantic[:, 3] = original[:, 3] + # Red minus green cancels the equal RGB renderer background, yielding + # the visible crown contribution, including bilinear/occlusion effects. + identity = outer if crown_uv is None else (crown_uv[group] | top) & outer + semantic[:, 0] = identity.float() + sl = slice(group * len(views), (group + 1) * len(views)) + expected = target[sl] * foreground[sl] + fg = foreground[sl] + + def objective(skins): + rendered = torch.stack([cropped.forward_view(skins, view) for view in views], 1) + probability = rendered[:, :, 0] - rendered[:, :, 1] + # The silhouette term also protects crown teeth above bare hair. + error = (probability - expected[None]).square() + error += (rendered[:, :, 3] - fg[None]).square() + return error.sum((1, 2, 3)) + + score = float(objective(semantic)) + record = {'before_loss': score, 'removed': []} + while True: + active = (top & (semantic[0, 3] > .5)).nonzero() + if not len(active): + break + scores = [] + for chunk in active.split(4): + batch = semantic.repeat(len(chunk), 1, 1, 1) + batch[torch.arange(len(chunk), device=uv.device), 3, chunk[:, 0], chunk[:, 1]] = 0 + scores.extend(objective(batch).tolist()) + best = min(range(len(scores)), key=scores.__getitem__) + # A positive numerical margin prevents deleting unobserved cells + # because of summation noise. It is not an occupancy prior. + if scores[best] >= score - .05: + break + y, x = active[best].tolist() + record['removed'].append({'uv': [x, y], 'loss_improvement': score - scores[best]}) + semantic[0, 3, y, x] = 0 + score = scores[best] + record['after_loss'] = score + records.append(record) + result[group, 3] = semantic[0, 3] + return result, records + + +@torch.no_grad() +def reconcile_crown_geometry(conditioning, details, renderer, views): + from SkingToolkit.dense_uv_parser.headwear import headwear_uv_families + from SkingToolkit.dense_uv_parser.infer import simple_inpaint_uv + routing, outputs = details['routing'], details['outputs'] + required = ('headwear_presence_logits', 'headwear_logits') + if len(views) < 2 or any(key not in outputs for key in required) or 'headwear_supported' not in routing: + return conditioning + royal = routing['headwear_supported'] & (routing['headwear_family'] == 4) + active = royal.flatten(1).any(1).reshape(-1, len(views)).all(1) + if not active.any(): + return conditioning + # Use the same completed base alpha that the final renderer will see. + # Material is discarded before comparison; RGB cannot influence geometry. + uv = torch.stack([simple_inpaint_uv(c[None].cpu())[0] for c in conditioning]).to(conditioning.device) + families = headwear_uv_families(routing, uv.shape[0], len(views)) + target = outputs['headwear_logits'].float().softmax(1)[:, 6] + changed = torch.zeros_like(uv[:, 3], dtype=torch.bool) + records = [] + for group in active.nonzero().flatten().tolist(): + sl = slice(group * len(views), (group + 1) * len(views)) + fixed, record = prune_crown_top(uv[group:group + 1], target[sl], routing['observed_foreground'][sl], + renderer, views, crown_uv=families[group:group + 1] == 4) + changed[group] = (uv[group, 3] > .5) & (fixed[0, 3] <= .5) + records.append({'group': group, **record[0]}) + result = conditioning.clone() + offset = 6 if conditioning.shape[1] == 12 else 5 + result[:, offset:] = result[:, offset:].masked_fill(changed[:, None], 0) + details['headwear_removed_top_uv'] = changed + details['crown_geometry'] = records + return result diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground.py new file mode 100644 index 0000000000000000000000000000000000000000..4e8a893fdb994596a5d2f43fec51a579ee1df0b7 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground.py @@ -0,0 +1,160 @@ +from pathlib import Path + +import torch +import torch.nn.functional as F +from torchvision.utils import save_image + + +ADAPTIVE_BACKGROUND_CANDIDATES = ( + (10, 10, 10), + (245, 245, 245), + (13, 204, 242), + (242, 26, 191), + (26, 230, 64), + (242, 191, 13), + (38, 64, 242), + (242, 51, 26), +) + + +def select_adaptive_background( + rendered, + foreground_mask, + boundary_width=2, + distance_quantile=0.10, +): + """Choose a deterministic parser background far from foreground edges.""" + if foreground_mask.dim() == 3: + foreground_mask = foreground_mask.unsqueeze(1) + expected = rendered.shape[:1] + (1,) + rendered.shape[-2:] + if foreground_mask.shape != expected: + raise ValueError( + f"Expected foreground mask shape {expected}, got " + f"{tuple(foreground_mask.shape)}." + ) + if boundary_width < 1: + raise ValueError("boundary_width must be positive.") + if not 0.0 <= distance_quantile <= 1.0: + raise ValueError("distance_quantile must be in [0, 1].") + + mask = foreground_mask.to(device=rendered.device, dtype=torch.bool) + kernel = boundary_width * 2 + 1 + eroded = -F.max_pool2d( + -mask.float(), kernel, stride=1, padding=boundary_width + ) + boundary = mask & (eroded < 0.5) + candidates = rendered.new_tensor(ADAPTIVE_BACKGROUND_CANDIDATES) / 255.0 + selected = [] + selected_indices = [] + for item in range(rendered.shape[0]): + sample_mask = boundary[item, 0] + if not sample_mask.any(): + sample_mask = mask[item, 0] + if not sample_mask.any(): + selected_indices.append(0) + selected.append(candidates[0]) + continue + colors = rendered[item, :3, sample_mask].transpose(0, 1).float() + distances = ( + colors.unsqueeze(0) - candidates.float().unsqueeze(1) + ).square().mean(dim=2).sqrt() + scores = torch.quantile(distances, distance_quantile, dim=1) + index = int(scores.argmax().item()) + selected_indices.append(index) + selected.append(candidates[index]) + return torch.stack(selected, dim=0).view(-1, 3, 1, 1), selected_indices + + +def build_parser_input( + rendered, + foreground_mask, + bg_color=(128, 128, 128), + background_mode="adaptive", + return_background=False, +): + """Composite a foreground mask into opaque RGB for the dense parser.""" + if foreground_mask.dim() == 3: + foreground_mask = foreground_mask.unsqueeze(1) + expected = rendered.shape[:1] + (1,) + rendered.shape[-2:] + if foreground_mask.shape != expected: + raise ValueError( + f"Expected foreground mask shape {expected}, got " + f"{tuple(foreground_mask.shape)}." + ) + mask = foreground_mask.to(device=rendered.device, dtype=torch.bool) + if background_mode == "adaptive": + background, selected_indices = select_adaptive_background( + rendered, mask + ) + elif background_mode == "neutral": + background = rendered.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 + background = background.expand(rendered.shape[0], -1, -1, -1) + selected_indices = None + else: + raise ValueError( + f"Unknown foreground parser background mode {background_mode!r}." + ) + + rgb = torch.where( + mask.expand(-1, 3, -1, -1), rendered[:, :3], background + ) + parser_input = torch.cat([rgb, torch.ones_like(rendered[:, 3:4])], dim=1) + if return_background: + return parser_input, background[:, :, 0, 0], selected_indices + return parser_input + + +def save_flood_outputs( + rendered, + foreground_mask, + view_count, + probability_output=None, + raw_mask_output=None, + mask_output=None, + cutout_output=None, +): + """Save the deterministic flood mask and transparent RGBA cutout.""" + if foreground_mask.dim() == 3: + foreground_mask = foreground_mask.unsqueeze(1) + expected = rendered.shape[:1] + (1,) + rendered.shape[-2:] + if foreground_mask.shape != expected: + raise ValueError( + f"Expected foreground mask shape {expected}, got " + f"{tuple(foreground_mask.shape)}." + ) + mask = foreground_mask.to(device=rendered.device, dtype=torch.bool) + score = mask.to(dtype=rendered.dtype) + cutout = torch.cat([rendered[:, :3], score], dim=1) + for tensor, output in ( + (score, probability_output), + (score, raw_mask_output), + (score, mask_output), + (cutout, cutout_output), + ): + if output is None: + continue + path = Path(output) + path.parent.mkdir(parents=True, exist_ok=True) + save_image(tensor.detach().cpu(), path, nrow=view_count) + return mask[:, 0] + + +def learned_foreground_support(probability, foreground_threshold=.5, source_threshold=.98, source_inset=1): + """Separate object occupancy from confident interior RGB sources. + + Segmentation confidence is not a calibrated alpha matte. Uncertain edge + colours are excluded from UV sampling rather than algebraically unmixed. + """ + if probability.dim()!=3 or not torch.isfinite(probability).all(): + raise ValueError('Expected finite foreground probability (N,H,W)') + if probability.min()<0 or probability.max()>1: + raise ValueError('Foreground probability must lie in [0,1]') + if not 0<=foreground_threshold<=source_threshold<=1 or source_inset<0: + raise ValueError('Invalid foreground/source thresholds or inset') + foreground=probability>=foreground_threshold + interior=foreground + for _ in range(source_inset): + padded=F.pad(interior[:,None].float(),(1,1,1,1),value=0) + interior=(-F.max_pool2d(-padded,3,1)[:,0])>.5 + sources=interior&(probability>=source_threshold) + return foreground,sources diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_acceptance.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_acceptance.py new file mode 100644 index 0000000000000000000000000000000000000000..a5fedc4e6fc5b35b3869c35d4211fd9398ef7275 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_acceptance.py @@ -0,0 +1,59 @@ +"""Audit frozen foreground predictions, including disconnected real character details.""" +import argparse +import hashlib +import json +from pathlib import Path +import numpy as np +from PIL import Image, ImageFilter + + +def inspect_real(directory): + root = Path(__file__).parent + rows = [] + for manifest_path in sorted(Path(directory).glob('*/manifest.json')): + manifest = json.loads(manifest_path.read_text()) + source = Path(manifest['input']) + if hashlib.sha256(source.read_bytes()).hexdigest() != manifest['input_sha256']: + raise ValueError('Regression input changed: ' + str(source)) + probability = np.array(Image.open(manifest_path.parent/'probability.png'))/255. + row = {'input': str(source), 'foreground_sha256': manifest.get('adaptation_sha256')} + if source.stem in {'1H6GAQA9BENWNB89_edited','2RSZETA9PCRZHVEU_edited','TWRLRRTHQP2UV368_edited','PU418PLGERMFC487_edited','img3','img23','img27','img28'}: + # Inspected empty background above the left sloping head edge. + row['annotated_background_false_positive_rate'] = [ + float((v[23:28,160:175]>=.5).mean()) for v in np.split(probability,2,axis=1)[:(2 if source.stem in {'TWRLRRTHQP2UV368_edited','PU418PLGERMFC487_edited','img3'} else 1)]] + if 'img28' == source.stem: + # Manually inspected interior rectangles at original 1024 x 1024. + # These are true disconnected blue character details, not background. + if probability.shape != (1024,1024): + raise ValueError('img28 annotation requires original dimensions') + boxes = [(262,860,280,880), (216,897,232,912), + (774,860,792,880), (728,897,744,912)] + row['disconnected_detail_recall'] = [ + float((probability[y1:y2,x1:x2]>=.5).mean()) for x1,y1,x2,y2 in boxes] + if 'PU418' in source.stem or 'TWRLR' in source.stem: + views = np.split(probability,2,axis=1) + native = [np.array(Image.fromarray(v.astype('float32')).resize((256,512),Image.Resampling.NEAREST)) for v in views] + if 'PU418' in source.stem: + mask = np.array(Image.open(root/'runs/v101_hat_audit_20260906/hat_brim_development_mask.png').filter(ImageFilter.MinFilter(3)))>0 + row['hat_brim_recall'] = [float((v[mask]>=.5).mean()) for v in native] + else: + mask = np.array(Image.open(root/'regression/TWRLRRTHQP2UV368_glasses.png').filter(ImageFilter.MinFilter(5)))>0 + row['glasses_recall'] = float((native[0][mask]>=.5).mean()) + rows.append(row) + if not rows: + raise ValueError('No prediction manifests') + return {'role':'real development annotations; never gradient training or held-out test', 'cases': rows} + + +def main(): + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--probability-dir', type=Path, required=True) + p.add_argument('--output', type=Path, required=True) + opt=p.parse_args() + result=inspect_real(opt.probability_dir) + opt.output.write_text(json.dumps(result,indent=2)) + print(json.dumps(result,indent=2)) + + +if __name__=='__main__': + main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_batch.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..b20dd8efc5225367d71c734cca2f20aefa5ba3c7 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_batch.py @@ -0,0 +1,53 @@ +"""Run the trained foreground model and UV parser with their existing environments.""" +import argparse +import json +import subprocess +import tempfile +import sys +from pathlib import Path + + +def main(): + root = Path(__file__).resolve().parents[1] + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--checkpoint', type=Path, required=True) + p.add_argument('--foreground-checkpoint', type=Path, required=True) + p.add_argument('--foreground-model-dir', type=Path, required=True) + p.add_argument('--foreground-python', default=sys.executable) + p.add_argument('--parser-python', default=sys.executable) + p.add_argument('--output-dir', type=Path, required=True) + p.add_argument('--probability-dir', type=Path, required=True) + p.add_argument('--pipeline', type=Path) + p.add_argument('--inputs', type=Path, nargs='+') + p.add_argument('--save-tensors', action='store_true') + opt = p.parse_args() + inputs = opt.inputs or [Path(x) for x in json.loads((root/'dense_uv_parser/v101_cases.json').read_text())] + for path in inputs + [opt.checkpoint, opt.foreground_checkpoint]: + if not path.is_file(): + raise FileNotFoundError(path) + entry = root/'dense_uv_parser/run_local.py' + with tempfile.TemporaryDirectory(prefix='foreground-batch-') as tmp: + manifest = Path(tmp)/'results.json' + subprocess.run([opt.foreground_python, str(entry), 'foreground_provider', + '--checkpoint', str(opt.foreground_checkpoint.resolve()), + '--model-dir', str(opt.foreground_model_dir.resolve()), + '--output-dir', str(opt.probability_dir.resolve()), + '--result-manifest', str(manifest), + '--inputs', *[str(x.resolve()) for x in inputs]], cwd=root, check=True) + probability = json.loads(manifest.read_text())['probabilities'] + if len(probability) != len(inputs): + raise ValueError('Foreground output count does not match inputs') + command = [opt.parser_python, str(entry), 'batch_accessories', + '--checkpoint', str(opt.checkpoint.resolve()), + '--output-dir', str(opt.output_dir.resolve()), + '--foreground-probabilities', *probability, + '--inputs', *[str(x.resolve()) for x in inputs]] + if opt.pipeline: + command += ['--pipeline', str(opt.pipeline.resolve())] + if opt.save_tensors: + command += ['--save-tensors'] + subprocess.run(command, cwd=root, check=True) + + +if __name__ == '__main__': + main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_model.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_model.py new file mode 100644 index 0000000000000000000000000000000000000000..1b861ff4ef62ce34a571682b78257bd734950bcf --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_model.py @@ -0,0 +1,32 @@ +"""Local pretrained BiRefNet plus versioned Minecraft decoder checkpoints.""" +import hashlib,io,json +from pathlib import Path +import torch + + +def sha(path): + with Path(path).open('rb') as f:return hashlib.file_digest(f,'sha256').hexdigest() + + +def load_foreground_model(model_dir,checkpoint=None,device='cuda'): + from transformers import AutoConfig,AutoModelForImageSegmentation + from safetensors.torch import load_file + path=Path(model_dir) + config=AutoConfig.from_pretrained(str(path),trust_remote_code=True,local_files_only=True) + model=AutoModelForImageSegmentation.from_config(config,trust_remote_code=True) + model.load_state_dict(load_file(path/'model.safetensors'),strict=True) + metadata=None + if checkpoint: + checkpoint_bytes=Path(checkpoint).read_bytes() + payload=torch.load(io.BytesIO(checkpoint_bytes),map_location='cpu',weights_only=False) + metadata={**payload['manifest'],'checkpoint_sha256':hashlib.sha256(checkpoint_bytes).hexdigest()} + if metadata['base_sha256']!=sha(path/'model.safetensors'):raise ValueError('Foreground base checkpoint mismatch') + missing,extra=model.load_state_dict(payload['decoder'],strict=False) + if extra or any(not name.startswith('bb.') for name in missing):raise ValueError((missing,extra)) + return model.eval().to(device),metadata + + +def foreground_logits(model,rgb): + mean=rgb.new_tensor([.485,.456,.406])[None,:,None,None] + std=rgb.new_tensor([.229,.224,.225])[None,:,None,None] + return model((rgb-mean)/std)[-1] diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_provider.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_provider.py new file mode 100644 index 0000000000000000000000000000000000000000..ddc0bfc379c08fe80576057c146433b100415d3e --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/foreground_provider.py @@ -0,0 +1,89 @@ +"""Offline BiRefNet foreground confidence; kept separate from UV/parser semantics.""" +import argparse +import hashlib +import io +import json +import os +import tempfile +from pathlib import Path + + +def sha(path): + with Path(path).open('rb') as f:return hashlib.file_digest(f,'sha256').hexdigest() + + +def main(): + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('--inputs',type=Path,nargs='+',required=True) + parser.add_argument('--output-dir',type=Path,required=True) + parser.add_argument('--model-dir',type=Path,required=True,help='Reviewed local BiRefNet snapshot with code and model.safetensors') + parser.add_argument('--result-manifest',type=Path,help='Write the exact ordered probability paths for a downstream batch') + parser.add_argument('--view-count',type=int,default=2) + parser.add_argument('--checkpoint',type=Path,help='Trained Minecraft foreground decoder checkpoint') + opt=parser.parse_args() + if opt.view_count<1:parser.error('View count must be positive') + os.environ['HF_HUB_OFFLINE']='1' + os.environ['HF_HUB_DISABLE_PROGRESS_BARS']='1' + import torch + import torch.nn.functional as F + import numpy as np + from PIL import Image + if __package__: + from .foreground_model import load_foreground_model + else: + from foreground_model import load_foreground_model + torch.set_num_threads(4) + model_path=opt.model_dir.resolve() + model_sources={p.name:sha(p) for p in model_path.glob('*.py')} + provenance={'model_dir':str(model_path),'model_sha256':sha(model_path/'model.safetensors'), + 'model_code_sha256':model_sources,'config_sha256':sha(model_path/'config.json'), + 'provider_sha256':sha(__file__),'loader_sha256':sha(Path(__file__).with_name('foreground_model.py')),'view_count':opt.view_count,'resolution':1024, + 'output':'Segmentation confidence, not a calibrated alpha matte','torch':torch.__version__} + model,trained_metadata=load_foreground_model(model_path,opt.checkpoint) + if opt.checkpoint: + provenance['adaptation_checkpoint']=str(opt.checkpoint.resolve()) + provenance['adaptation_sha256']=trained_metadata['checkpoint_sha256'] + provenance['adaptation_version']=trained_metadata['version'] + provenance['output']='Supervised foreground alpha prediction with separate reliable RGB sampling' + mean=torch.tensor([.485,.456,.406],device='cuda')[None,:,None,None] + std=torch.tensor([.229,.224,.225],device='cuda')[None,:,None,None] + opt.output_dir.mkdir(parents=True,exist_ok=True) + results=[] + for path in opt.inputs: + input_bytes=path.read_bytes() + manifest={**provenance,'input':str(path.resolve()),'input_sha256':hashlib.sha256(input_bytes).hexdigest()} + fingerprint=hashlib.sha256(json.dumps(manifest,sort_keys=True).encode()).hexdigest() + target=opt.output_dir/(path.stem+'_'+fingerprint[:12]) + if (target/'manifest.json').is_file(): + prior=json.loads((target/'manifest.json').read_text()) + if prior.get('fingerprint')==fingerprint and (target/'probability.png').is_file(): + results.append(str((target/'probability.png').resolve())) + print('cached='+str(target/'probability.png'),flush=True);continue + if target.exists():raise FileExistsError(target) + im=np.asarray(Image.open(io.BytesIO(input_bytes)).convert('RGB')).copy() + if im.shape[1]%opt.view_count:raise ValueError('Image width must divide evenly into views') + rgb=torch.from_numpy(im).permute(2,0,1).float().cuda()/255 + views=torch.stack(rgb.chunk(opt.view_count,dim=2)) + x=F.interpolate(views,size=(1024,1024),mode='bilinear',align_corners=False) + with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16): + logits=model((x-mean)/std)[-1] + confidence=F.interpolate(logits.float().sigmoid(),views.shape[-2:],mode='bilinear',align_corners=False)[:,0] + if not torch.isfinite(confidence).all():raise RuntimeError('Non-finite foreground prediction') + combined=torch.cat(list(confidence),dim=1).cpu().numpy() + # Exact grayscale output is hashed again by batch_accessories. + with tempfile.TemporaryDirectory(prefix='.foreground-',dir=opt.output_dir) as temporary: + scratch=Path(temporary) + Image.fromarray(np.uint8(combined*255)).save(scratch/'probability.png') + manifest['fingerprint']=fingerprint + (scratch/'manifest.json').write_text(json.dumps(manifest,indent=2)) + os.replace(scratch,target) + results.append(str((target/'probability.png').resolve())) + print('completed='+str(target/'probability.png'),flush=True) + if opt.result_manifest: + opt.result_manifest.parent.mkdir(parents=True,exist_ok=True) + temporary=opt.result_manifest.with_suffix('.tmp') + temporary.write_text(json.dumps({'probabilities':results,'provenance':provenance},indent=2)) + os.replace(temporary,opt.result_manifest) + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/head_inpainting.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/head_inpainting.py new file mode 100644 index 0000000000000000000000000000000000000000..c6995db7ae246ca3b07fdb759365f021690deb82 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/head_inpainting.py @@ -0,0 +1,43 @@ +"""Use observed hair as the source for hidden hair beside an accepted accessory.""" +import torch +from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology + + +@torch.no_grad() +def repair_hidden_head_material(uv,conditioning,details,views): + routing=details['routing'];rejected=routing.get('headphone_colour_rejected') + logits=details['outputs'].get('head_ownership_logits') + if rejected is None or logits is None or not rejected.any():return uv + topology=build_simple_uv_topology();device=uv.device + head=(topology.valid&(topology.part==0)&(topology.layer==0)).to(device).flatten() + face=topology.face.to(device).flatten();position=topology.world_position.to(device).reshape(-1,3) + groups=uv.shape[0];count=uv.new_zeros(groups,4096);hair=uv.new_zeros(groups,4096) + probability=logits.float().softmax(1)[:,2] + for vi in range(len(views)): + sl=slice(vi,routing['layer'].shape[0],len(views)) + valid=routing['color_foreground'][sl]&(routing['layer'][sl]==0)&(routing['part'][sl]==0) + index=routing['flat_uv'][sl].flatten(1) + count.scatter_add_(1,index,valid.float().flatten(1)) + hair.scatter_add_(1,index,(probability[sl]*valid).flatten(1)) + hair=hair/count.clamp_min(1) + result=uv.clone().flatten(2);known=(conditioning[:,4]>.5).flatten(1) + for batch in range(groups): + if not rejected[batch*len(views):(batch+1)*len(views)].any():continue + reliable=known[batch]&head&(count[batch]>=8)&(hair[batch]>=.8) + for side in (1,2,3,4): + observed=known[batch]&head&(face==side)&(count[batch]>=8) + sources=reliable&(face==side) + if int(sources.sum())<2: + # With a largely occluded ear-side face, use the observed hair on + # adjacent head faces, provided no visible face evidence contradicts it. + if int(observed.sum())>0 and float(sources.sum()/observed.sum())<.6:continue + sources=reliable + if int(sources.sum())<2:continue + if observed.any() and float((reliable&observed).sum()/observed.sum())<.6:continue + targets=(~known[batch]&head&(face==side)).nonzero().flatten() + if not len(targets):continue + source_indices=sources.nonzero().flatten() + nearest=torch.cdist(position[targets],position[source_indices]).argmin(1) + # Never recursively sample a repaired texel or change an observed one. + result[batch,:3,targets]=uv.flatten(2)[batch,:3,source_indices[nearest]] + return result.reshape_as(uv) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear.py new file mode 100644 index 0000000000000000000000000000000000000000..04fba716ca198161ff097a4fb9921a2d37676daf --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear.py @@ -0,0 +1,206 @@ +"""Semantic headwear components with one layer decision per player/component. + +The crown class denotes a royal crown, distinct from the body of a fabric hat. +Colours are never used to create identities or decide their layer. +""" +import torch +from torch import nn +from SkingToolkit.dense_uv_parser.accessories import HeadAccessoryHead, connected_object_support + +HEADWEAR_CLASSES = ('none', 'inner_hat_body', 'inner_hat_band', + 'outer_hat_body', 'outer_hat_band', 'brim', 'royal_crown') + + +class HeadwearHead(HeadAccessoryHead): + def __init__(self, semantic_dim=768, predict_presence=False): + super().__init__(semantic_dim, False) + self.classifier = nn.Conv2d(24, len(HEADWEAR_CLASSES), 1) + dim = 7*8*8+semantic_dim + self.presence = nn.Sequential(nn.LayerNorm(dim),nn.Linear(dim,192),nn.SiLU(),nn.Dropout(.1),nn.Linear(192,2)) if predict_presence else None + + def forward(self,crop,semantic_features,return_presence=False): + logits=super().forward(crop,semantic_features) + if not return_presence:return logits + spatial=torch.nn.functional.adaptive_avg_pool2d(logits.float().softmax(1),(8,8)).flatten(1) + semantic=semantic_features.float().mean((2,3)) + return logits,self.presence(torch.cat([spatial,semantic],1)) + + +def component_decisions(logits, foreground, views_per_group, seed_threshold=.80): + """Pool both views BEFORE routing. A band cannot vote twice for two layers. + + A single player has one hat body, one decorative band, and one brim/crown. + Separate players are never pooled. Background and low confidence isolated + regions cannot establish a component. Holes stay outside its learned mask. + """ + if logits.shape[0] % views_per_group: + raise ValueError('Incomplete headwear view group') + p = logits.float().softmax(1) + family_p = torch.stack([p[:,0], p[:,1]+p[:,3], p[:,2]+p[:,4], p[:,5], p[:,6]], 1) + confidence, family = family_p.max(1) + supported = torch.zeros_like(foreground) + layer = torch.full_like(family, -1) + for component in (1,2,3,4): + candidates = (family == component) & foreground + grown, _ = connected_object_support(confidence, candidates.long(), foreground, + seed_threshold, .5) + for start in range(0, logits.shape[0], views_per_group): + sl = slice(start, start+views_per_group) + mask = grown[sl] + if not mask.any(): + continue + # A few confident speckles are insufficient to claim an object. + if (confidence[sl][mask] >= seed_threshold).sum() < 16: + continue + selected_layer = 1 + if component in (1,2): + inner_id, outer_id = (1,3) if component == 1 else (2,4) + inner_evidence = p[sl,inner_id][mask].sum() + outer_evidence = p[sl,outer_id][mask].sum() + selected_layer = int(outer_evidence > inner_evidence) + supported[sl] |= mask + layer[sl][mask] = selected_layer + return supported, layer, family, confidence + + +def apply_headwear_routing(routing, outputs, foreground, renderer, views): + from SkingToolkit.dense_uv_parser.utils import build_static_surface_routing + logits = outputs.get('headwear_logits') + if logits is None: + return + support, layer, family, confidence = component_decisions(logits, foreground, len(views)) + presence=outputs.get('headwear_presence_logits') + if presence is not None: + # Object classification uses both semantic shape and pretrained image + # features. Require agreement of the two views before changing layers. + probability=presence.sigmoid().reshape(-1,len(views),2).amin(1).repeat_interleave(len(views),0) + hat=probability[:,0]>=outputs.get('headwear_presence_threshold',.95) + crown=probability[:,1]>=outputs.get('headwear_presence_threshold',.95) + support &= torch.where(family==4,crown[:,None,None],hat[:,None,None]) + layer=layer.masked_fill(~support,-1) + semantic_support = support.clone() + # Resolve residual old-route pixels at the output texel resolution. A thin + # uncertain fringe cannot retain an outer cell whose footprint is mostly an + # explicitly inner component. This uses semantic votes, never colour area. + groups = foreground.shape[0] // len(views) + votes = confidence.new_zeros(groups,4096) + counts = torch.zeros_like(votes) + for vi,view in enumerate(views): + sl=slice(vi,foreground.shape[0],len(views)) + static=build_static_surface_routing(renderer,view,foreground.device) + valid=foreground[sl]&static['masks'][1]&(static['part'][1]==0) + index=static['flat_uv'][1].flatten()[None].expand(groups,-1) + votes.scatter_add_(1,index,(valid&support[sl]&(layer[sl]==0)).float().flatten(1)) + counts.scatter_add_(1,index,valid.float().flatten(1)) + inner_cells=(votes/counts.clamp_min(1)>=.6)&(counts>=8) + veto=torch.zeros_like(support) + for vi,view in enumerate(views): + sl=slice(vi,foreground.shape[0],len(views)) + static=build_static_surface_routing(renderer,view,foreground.device) + cell=inner_cells.gather(1,static['flat_uv'][1].flatten()[None].expand(groups,-1)).reshape_as(support[sl]) + accepted=cell&foreground[sl]&static['masks'][0]&(static['part'][0]==0)&~(support[sl]&(layer[sl]==1)) + accepted &= ~(((family[sl]==3)|(family[sl]==4)) & (confidence[sl]>=.5)) + veto[sl]=accepted;support[sl]|=accepted;layer[sl][accepted]=0 + accepted_all = torch.zeros_like(support) + for vi, view in enumerate(views): + sl = slice(vi, foreground.shape[0], len(views)) + static = build_static_surface_routing(renderer, view, foreground.device) + for selected in (0,1): + accepted = support[sl] & (layer[sl] == selected) & static['masks'][selected] & (static['part'][selected] == 0) + accepted_all[sl] |= accepted + for name, value in [('layer',selected),('foreground',True),('surface',selected), + ('secondary',False),('secondary_routed',False), + ('semantic_fallback',False),('consensus_outer_gate_rejected',False)]: + if name in routing: + routing[name][sl] = torch.where(accepted, torch.full_like(routing[name][sl],value), routing[name][sl]) + for name in ('flat_uv','part','face','texel_center_score'): + if name in routing: + routing[name][sl] = torch.where(accepted, static[name][selected], routing[name][sl]) + routing['confidence'][sl] = torch.where(accepted, confidence[sl], routing['confidence'][sl]) + margin = (2*confidence[sl]-1).clamp_min(0) + for name, value in [('confidence_margin',margin),('confidence_margin_ratio',margin/confidence[sl].clamp_min(1e-6))]: + if name in routing: + routing[name][sl] = torch.where(accepted,value,routing[name][sl]) + routing['accessory_supported'] = (routing['accessory_supported'] | (accepted_all & (layer==1))) & ~(accepted_all & (layer==0)) + routing['headwear_supported'] = accepted_all + routing['headwear_layer'] = layer + routing['headwear_family'] = torch.where(accepted_all & semantic_support,family,0) + routing['headwear_probability'] = confidence + routing['headwear_inner_cell_veto'] = veto + + +def headwear_uv_families(routing, groups, views_per_group): + """Assign only observed headwear texels; none stays an explicit identity.""" + valid=routing['color_foreground'] & (routing['part']==0) + family=routing['headwear_family'] + confidence=routing['headwear_probability'] + votes=confidence.new_zeros(groups,5,4096) + for vi in range(views_per_group): + sl=slice(vi,valid.shape[0],views_per_group) + index=routing['flat_uv'][sl].flatten(1) + for cls in range(5): + evidence=valid[sl] & (family[sl]==cls) & (confidence[sl]>=.7) + votes[:,cls].scatter_add_(1,index,evidence.float().flatten(1)) + count,identity=votes.max(1) + return identity.masked_fill(count<4,0).reshape(groups,64,64) + + +def isolate_headwear_material(uv,conditioning,details,views): + """Unknown base texels never use a decorative band as an inpainting seed.""" + from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology + routing=details['routing'] + if 'headwear_supported' not in routing or not routing['headwear_supported'].any():return uv,None + family=headwear_uv_families(routing,uv.shape[0],len(views)) + topology=build_simple_uv_topology();device=uv.device + head=(topology.valid&(topology.part==0)&(topology.layer==0)).to(device).flatten() + faces=topology.face.to(device).flatten() + positions=topology.world_position.to(device).reshape(-1,3) + known=(conditioning[:,4]>.5).flatten(1) + original=uv.flatten(2);result=original.clone() + for group in range(uv.shape[0]): + band=(family[group].flatten()==2)&head + if not band.any():continue + for face in range(6): + target=(head&(faces==face)&~known[group]).nonzero().flatten() + source=(head&(faces==face)&known[group]&~band).nonzero().flatten() + if not target.numel() or not source.numel():continue + nearest=torch.cdist(positions[target],positions[source]).argmin(1) + result[group,:3,target]=original[group,:3,source[nearest]] + return result.reshape_as(uv),family + + +def reconcile_crown_top(conditioning,details,renderer,views): + """Keep exposed top hair when views disagree about crown occupancy. + + Edited views are not always a consistent 3D observation. An isolated crown + observation must not cover a clearly observed hair opening in another view. + No RGB comparison, fixed crown outline, or inner material transfer is used. + """ + from SkingToolkit.dense_uv_parser.utils import build_static_surface_routing + r=details['routing'];outputs=details['outputs'] + if len(views)<2 or 'headwear_presence_logits' not in outputs or 'headwear_supported' not in r or 'head_ownership_logits' not in outputs:return conditioning + royal=r['headwear_supported']&(r['headwear_family']==4) + if not royal.any():return conditioning + groups=conditioning.shape[0];device=conditioning.device + hair=outputs['head_ownership_logits'].float().softmax(1)[:,2] + no_headwear=outputs['headwear_logits'].float().softmax(1)[:,0] + seen=[];negative=[];positive=[] + for vi,view in enumerate(views): + sl=slice(vi,hair.shape[0],len(views)) + static=build_static_surface_routing(renderer,view,device) + valid=r['observed_foreground'][sl]&static['masks'][1]&(static['part'][1]==0)&(static['face'][1]==4) + index=static['flat_uv'][1].flatten()[None].expand(groups,-1) + count=conditioning.new_zeros(groups,4096);hn=torch.zeros_like(count);nn=torch.zeros_like(count);rp=torch.zeros_like(count) + count.scatter_add_(1,index,valid.float().flatten(1)) + hn.scatter_add_(1,index,(hair[sl]*valid).flatten(1));nn.scatter_add_(1,index,(no_headwear[sl]*valid).flatten(1)) + rp.scatter_add_(1,index,(royal[sl]&valid).float().flatten(1)) + observed=count>=16 + seen.append(observed);negative.append(observed&(hn/count.clamp_min(1)>.9)&(nn/count.clamp_min(1)>.9)) + positive.append(rp>=4) + conflict=torch.stack(seen).all(0)&torch.stack(negative).any(0)&torch.stack(positive).any(0) + conflict=conflict.reshape(groups,64,64) + if not conflict.any():return conditioning + result=conditioning.clone();offset=6 if conditioning.shape[1]==12 else 5 + result[:,offset:]=result[:,offset:].masked_fill(conflict[:,None],0) + details['headwear_removed_top_uv']=conflict + return result diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear_data.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear_data.py new file mode 100644 index 0000000000000000000000000000000000000000..06089e4e48a9cdb6dbecd7d0fa333de7797291d1 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/headwear_data.py @@ -0,0 +1,131 @@ +"""Authored headwear instances: whole bands and open, jewelled royal crowns.""" +import numpy as np +import torch +from torch.utils.data import Dataset +from SkingToolkit.dense_uv_parser.accessory_data import make_accessory_skin, render_accessories +from SkingToolkit.dense_uv_parser.ownership_data import make_ownership_skin +from SkingToolkit.dense_uv_parser.skin_dataset import load_skin + + +def make_headwear_skin(original, seed): + rng = np.random.default_rng(seed+807101) + if rng.random() < .28: + # Ear cups and a top band can resemble a crown in a head-only crop. + # Keep their complete authored geometry as explicit hard negatives. + for offset in range(100): + negative_uv, ownership = make_ownership_skin(original, seed+offset*7919) + if (ownership==4).any(): + return negative_uv, torch.zeros((64,64),dtype=torch.long) + raise RuntimeError("Failed to generate headphone negative") + uv, objects, components = make_accessory_skin(original, seed, return_components=True) + if rng.random() < .5: + return uv, components + # Start with explicitly authored negative face/hair/eyewear/headphones. A + # crown is a thin object around hair; it never labels the exposed hair top. + uv, _ = make_ownership_skin(original, seed+915103) + uv = uv.permute(1,2,0).numpy().copy() + uv[:16,32:] = 0 + components = np.zeros((64,64), dtype=np.int64) + # Some samples are difficult negatives with no headwear at all. + negative = rng.random() < .22 + faces = [(8,8),(24,8),(0,8),(16,8),(8,0)] + # Reauthor bare head: the negative generator may itself have had a hat. + hair = rng.uniform(.015,.40,3) if rng.random()<.7 else rng.uniform(.1,.9,3) + skin = rng.uniform(.25,.9,3) + for fi,(x,y) in enumerate(faces): + base = np.broadcast_to(hair if fi in (1,4) else skin,(8,8,3)).copy() + if fi in (0,2,3): base[:int(rng.integers(2,5))] = hair + uv[y:y+8,x:x+8,:3] = np.clip(base+rng.normal(0,.025,(8,8,1)),0,1) + uv[y:y+8,x:x+8,3] = 1 + uv[11:13,9:11,:3] = rng.uniform(0,.9,3) + uv[11:13,13:15,:3] = rng.uniform(0,.9,3) + if negative: + return torch.from_numpy(uv).permute(2,0,1).float(), torch.from_numpy(components) + yy,xx = np.mgrid[:8,:8] + metal = np.array([rng.uniform(.65,1),rng.uniform(.42,.85),rng.uniform(.02,.3)]) + if rng.random()<.45: metal = rng.uniform(.1,.95,3) # palette is not identity + band_y = int(rng.integers(1,4)); thickness = int(rng.integers(1,3)) + period = int(rng.integers(2,5)); phase = int(rng.integers(period)) + for fi,(x,y) in enumerate(faces): + x += 32 + if fi==4: + # Disconnected teeth at the top rim; the centre remains open. + mask=((yy==0)|(yy==7)|(xx==0)|(xx==7)) & (((xx+yy+phase)%period)==0) + if rng.random()<.5: mask[:]=False + else: + ring=(yy>=band_y)&(yy 0 + or model.semantic_spatial_feature_dim > 0 + ): + semantic_backbone = model_config.get("semantic_backbone", "siglip2") + semantic_model = model_config.get( + "semantic_model", + model_config.get( + "tipsv2_model" + if semantic_backbone == "tipsv2" + else "siglip_model", + checkpoint_args.get( + "tipsv2_model" + if semantic_backbone == "tipsv2" + else "siglip_model", + "google/tipsv2-b14" + if semantic_backbone == "tipsv2" + else "google/siglip2-base-patch16-224", + ), + ), + ) + attach_semantic_runtime( + model, + semantic_backbone, + semantic_model, + device, + local_files_only=bool( + checkpoint_args.get( + "tipsv2_local_files_only" + if semantic_backbone == "tipsv2" + else "siglip_local_files_only", + False, + ) + ), + runtime_batch_size=int( + checkpoint_args.get("semantic_runtime_batch_size", 32) + ), + ) + model.eval() + return model, checkpoint_args + + +def load_view_images(args, views, renderer, bg_color=(128, 128, 128)): + images = [] + if args.combined: + combined = Image.open(args.combined) + width, height = combined.size + if width % len(views) != 0: + raise ValueError(f"Combined image width {width} is not divisible by {len(views)} views.") + view_width = width // len(views) + images = [combined.crop((i * view_width, 0, (i + 1) * view_width, height)) for i in range(len(views))] + elif args.view_images: + if len(args.view_images) != len(views): + raise ValueError(f"Expected {len(views)} --view_images, got {len(args.view_images)}.") + images = [Image.open(path) for path in args.view_images] + elif args.front and args.back: + if len(views) != 2: + raise ValueError(f"--front/--back only works for 2-view checkpoints, got {len(views)} views.") + images = [Image.open(args.front), Image.open(args.back)] + else: + raise ValueError("Provide --combined, --view_images, or both --front and --back.") + + tensors = [ + image_to_render_tensor(image, view_native_size(renderer, view), bg_color=bg_color) + for image, view in zip(images, views) + ] + return torch.stack(tensors, dim=0) + + +def save_conditioning_preview(conditioning, output_path): + inner_rgb = conditioning[:, 0:3] + outer_offset = 6 if conditioning.shape[1] == 12 else 5 + outer_rgb = conditioning[:, outer_offset : outer_offset + 3] + preview = torch.cat([inner_rgb, outer_rgb], dim=0) + output_path.parent.mkdir(parents=True, exist_ok=True) + save_image(preview.clamp(0.0, 1.0), output_path, nrow=conditioning.shape[0]) + + +def save_parser_uv( + conditioning, + output_path, + alpha_threshold=0.5, + enforce_base_alpha=False, +): + """Save the partial parser atlas with unknown texels left transparent.""" + parser_uv = conditioning_to_pred_uv(conditioning) + if parser_uv.dim() != 4 or parser_uv.shape[0] != 1: + raise ValueError( + "Parser UV PNG output requires exactly one conditioning sample, " + f"got {tuple(parser_uv.shape)}." + ) + parser_uv = finalize_minecraft_alpha( + parser_uv[0], + alpha_threshold=alpha_threshold, + enforce_base_alpha=enforce_base_alpha, + ) + # A partial parser atlas is a diagnostic artifact, not a complete skin. + # Clear placeholder RGB under transparent texels so viewers that mishandle + # base-layer alpha cannot display the conditioning background as predicted + # skin color. + opaque = parser_uv[3:4] > 0.5 + parser_uv[:3] = torch.where( + opaque.expand_as(parser_uv[:3]), + parser_uv[:3], + torch.zeros_like(parser_uv[:3]), + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + tensor_to_rgba_image(parser_uv.detach().cpu()).save(output_path) + print( + "parser_uv_stats=" + + json.dumps( + { + "opaque_pixels": int(opaque.sum().item()), + "transparent_pixels": int((~opaque).sum().item()), + }, + sort_keys=True, + ) + ) + print(f"Saved parser_partial_uv={output_path}") + + +def simple_inpaint_uv(conditioning, alpha_threshold=0.5): + """Repair inner UV holes while preserving the parser's outer layer.""" + parser_uv = conditioning_to_pred_uv(conditioning) + if parser_uv.dim() != 4 or parser_uv.shape[0] != 1: + raise ValueError( + "Simple parser UV inpainting requires exactly one conditioning " + f"sample, got {tuple(parser_uv.shape)}." + ) + repaired, stats = simple_symmetry_nearest_inpaint( + parser_uv[0], alpha_threshold=alpha_threshold + ) + repaired = finalize_minecraft_alpha( + repaired, + alpha_threshold=alpha_threshold, + enforce_base_alpha=False, + ) + opaque = repaired[3:4] > 0.5 + repaired[:3] = torch.where( + opaque.expand_as(repaired[:3]), + repaired[:3], + torch.zeros_like(repaired[:3]), + ) + return repaired, stats + + +def save_simple_inpaint_uv(conditioning, output_path, alpha_threshold=0.5): + repaired, stats = simple_inpaint_uv( + conditioning, + alpha_threshold=alpha_threshold, + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + tensor_to_rgba_image(repaired.detach().cpu()).save(output_path) + print("simple_inpaint_stats=" + json.dumps(stats, sort_keys=True)) + print(f"Saved parser_simple_inpaint_uv={output_path}") + return repaired, stats + + +def _raw_debug_foreground( + outputs, + routing, + fg_threshold, + observed_foreground=None, +): + """Mask raw semantic heads with observed input foreground when available. + + The routed foreground is deliberately stricter: it excludes pixels rejected + by route confidence, margin, and UV coverage filters. Using it for a raw + head preview makes those routing rejections look like holes in the semantic + prediction, even though the head produced a class at every foreground pixel. + """ + model_foreground = ( + torch.sigmoid(outputs["foreground"])[:, 0] > fg_threshold + ) + if observed_foreground is not None: + if observed_foreground.shape != model_foreground.shape: + raise ValueError( + "observed_foreground must have shape " + f"{tuple(model_foreground.shape)}, got " + f"{tuple(observed_foreground.shape)}." + ) + return observed_foreground.to( + device=model_foreground.device, + dtype=torch.bool, + ) + if routing is None: + return model_foreground + return routing.get( + "observed_foreground", + routing.get("raw_foreground", model_foreground), + ) + + +def save_debug_preview( + rendered, + outputs, + view_count, + output_path, + fg_threshold, + bg_color=(128, 128, 128), + routing=None, + overlay_output=None, + overlay_alpha=0.45, + inner_cutout_output=None, + outer_cutout_output=None, + secondary_cutout_output=None, + color_source_output=None, + face_output=None, + layer_face_output=None, + raw_face_output=None, + raw_layer_face_output=None, + geometry_grid_output=None, + geometry_overlay_output=None, + geometry_routed_overlay_output=None, + geometry_fill_output=None, + renderer=None, + views=None, + raw_outputs=None, + raw_observed_foreground=None, +): + if not 0.0 <= overlay_alpha <= 1.0: + raise ValueError(f"overlay_alpha must be in [0, 1], got {overlay_alpha}.") + pred_fg = routing["foreground"] if routing is not None else torch.sigmoid(outputs["foreground"])[:, 0] > fg_threshold + pred_part_values = ( + outputs["part"].argmax(dim=1) + if "part" in outputs + else routing["part"] + ) + pred_part = torch.where(pred_fg, pred_part_values, torch.full_like(pred_part_values, IGNORE_INDEX)) + raw_head_outputs = outputs if raw_outputs is None else raw_outputs + raw_layer_values = raw_head_outputs["layer"].argmax(dim=1) + raw_fg = _raw_debug_foreground( + raw_head_outputs, + None if raw_outputs is not None else routing, + fg_threshold, + observed_foreground=raw_observed_foreground, + ) + pred_layer_values = routing["layer"] if routing is not None else raw_layer_values + pred_layer = torch.where( + pred_fg, + pred_layer_values, + torch.full_like(outputs["layer"].argmax(dim=1), IGNORE_INDEX), + ) + raw_face_values = ( + raw_head_outputs["face"].argmax(dim=1) + if "face" in raw_head_outputs + else routing["face"] + ) + raw_face = torch.where( + raw_fg, + raw_face_values, + torch.full_like(raw_face_values, IGNORE_INDEX), + ) + pred_face_values = routing["face"] if routing is not None else raw_face_values + pred_face = torch.where( + pred_fg, + pred_face_values, + torch.full_like(pred_face_values, IGNORE_INDEX), + ) + if "layer_face" in raw_head_outputs: + raw_layer_face_values = raw_head_outputs["layer_face"].argmax(dim=1) + raw_layer_face = torch.where( + raw_fg, + raw_layer_face_values, + torch.full_like(raw_layer_face_values, IGNORE_INDEX), + ) + else: + raw_layer_values_for_debug = raw_layer_values + if raw_head_outputs["layer"].shape[1] == 3: + # Geometry-fit class 2 is a secondary/backface role, not a skin + # layer. It has no raw inner/outer value until geometry routing, + # so leave it unlabelled in the pre-affine joint preview. + raw_layer_values_for_debug = raw_layer_values.masked_fill( + raw_layer_values == 2, + IGNORE_INDEX, + ) + raw_layer = torch.where( + raw_fg, + raw_layer_values_for_debug, + torch.full_like(raw_layer_values_for_debug, IGNORE_INDEX), + ) + raw_layer_face = combine_layer_face(raw_layer, raw_face) + pred_layer_face = combine_layer_face(pred_layer, pred_face) + pred_uv = flat_uv_to_uv01(routing["flat_uv"], rendered.dtype) if routing is not None else prediction_uv01(outputs) + + part_color = colorize_labels(pred_part, PART_PALETTE, bg_color, rendered) + layer_color = colorize_labels(pred_layer, LAYER_PALETTE, bg_color, rendered) + route_role_values = routing.get("route_role", raw_layer_values) if routing is not None else raw_layer_values + route_role_mask = ( + pred_fg | routing.get("secondary", torch.zeros_like(pred_fg)) + if routing is not None + else raw_fg + ) + route_role = torch.where( + route_role_mask, + route_role_values, + torch.full_like(route_role_values, IGNORE_INDEX), + ) + route_role_color = colorize_labels(route_role, ROUTE_ROLE_PALETTE, bg_color, rendered) + raw_face_color = colorize_labels(raw_face, FACE_PALETTE, bg_color, rendered) + face_color = colorize_labels(pred_face, FACE_PALETTE, bg_color, rendered) + raw_layer_face_color = colorize_labels(raw_layer_face, LAYER_FACE_PALETTE, bg_color, rendered) + layer_face_color = colorize_labels(pred_layer_face, LAYER_FACE_PALETTE, bg_color, rendered) + geometry_images = None + geometry_overlays = None + geometry_routed_overlays = None + if renderer is not None and views is not None and routing is not None: + geometry_debug = build_geometry_grid_debug( + renderer, views, rendered.shape[0], rendered, bg_color=bg_color + ) + inner_grid, outer_grid = geometry_debug[:2] + inner_fill, outer_fill = fill_geometry_grid_debug( + rendered, pred_fg, pred_layer_values, geometry_debug, bg_color=bg_color + ) + geometry_images = (inner_grid, outer_grid, inner_fill, outer_fill) + geometry_overlays = overlay_geometry_grid_debug(rendered, geometry_debug) + geometry_routed_overlays = overlay_geometry_grid_debug( + rendered, + geometry_debug, + base_images=(inner_fill, outer_fill), + ) + debug_images = [ + rendered[:, :3], + colorize_foreground(pred_fg, bg_color, rendered), + part_color, + layer_color, + route_role_color, + raw_face_color, + face_color, + raw_layer_face_color, + layer_face_color, + ] + if geometry_images is not None: + debug_images.extend(geometry_overlays) + debug_images.extend(geometry_routed_overlays) + debug_images.extend(geometry_images) + surface_color = None + if routing is not None: + pred_surface = torch.where( + pred_fg, + routing["surface"], + torch.full_like(routing["surface"], IGNORE_INDEX), + ) + surface_color = colorize_surface(pred_surface, bg_color, rendered) + debug_images.append(surface_color) + debug_images.append(colorize_uv(pred_uv, pred_fg, bg_color)) + if output_path is not None: + debug_preview = torch.cat(debug_images, dim=0) + output_path.parent.mkdir(parents=True, exist_ok=True) + save_image(debug_preview.clamp(0.0, 1.0).detach().cpu(), output_path, nrow=view_count) + + for colorized, path in ( + (face_color, face_output), + (layer_face_color, layer_face_output), + (raw_face_color, raw_face_output), + (raw_layer_face_color, raw_layer_face_output), + ): + if path is not None: + path.parent.mkdir(parents=True, exist_ok=True) + save_image(colorized.clamp(0.0, 1.0).detach().cpu(), path, nrow=view_count) + + if geometry_images is not None: + for images, path in ( + (geometry_overlays, geometry_overlay_output), + (geometry_routed_overlays, geometry_routed_overlay_output), + (geometry_images[:2], geometry_grid_output), + (geometry_images[2:], geometry_fill_output), + ): + if path is not None: + path.parent.mkdir(parents=True, exist_ok=True) + save_image( + torch.cat(images, dim=0).clamp(0.0, 1.0).detach().cpu(), + path, + nrow=view_count, + ) + + if ( + overlay_output is not None + or inner_cutout_output is not None + or outer_cutout_output is not None + or secondary_cutout_output is not None + or color_source_output is not None + ): + rgb = rendered[:, :3] + mask = pred_fg.unsqueeze(1) + bg = rgb.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 + routed_original = torch.where(mask, rgb, bg.expand_as(rgb)) + inner_mask = (pred_fg & (pred_layer_values == 0)).unsqueeze(1) + outer_mask = (pred_fg & (pred_layer_values == 1)).unsqueeze(1) + secondary_mask = ( + routing.get("secondary", raw_fg & (raw_layer_values == 2)) + if routing is not None + else raw_fg & (raw_layer_values == 2) + ).unsqueeze(1) + inner_cutout = torch.where(inner_mask, rgb, bg.expand_as(rgb)) + outer_cutout = torch.where(outer_mask, rgb, bg.expand_as(rgb)) + secondary_cutout = torch.where(secondary_mask, rgb, bg.expand_as(rgb)) + color_source_mask = ( + routing.get("color_foreground", pred_fg) + if routing is not None + else pred_fg + ).unsqueeze(1) + color_source_cutout = torch.where( + color_source_mask, rgb, bg.expand_as(rgb) + ) + + def save_cutout(cutout, path): + if path is None: + return + path.parent.mkdir(parents=True, exist_ok=True) + save_image(cutout.clamp(0.0, 1.0).detach().cpu(), path, nrow=view_count) + + save_cutout(inner_cutout, inner_cutout_output) + save_cutout(outer_cutout, outer_cutout_output) + save_cutout(secondary_cutout, secondary_cutout_output) + save_cutout(color_source_cutout, color_source_output) + + def overlay(colorized): + blended = rgb * (1.0 - overlay_alpha) + colorized * overlay_alpha + return torch.where(mask, blended, rgb) + + if overlay_output is not None: + overlay_images = [ + rgb, + routed_original, + *(geometry_overlays or ()), + *(geometry_routed_overlays or ()), + inner_cutout, + outer_cutout, + secondary_cutout, + overlay(part_color), + overlay(layer_color), + overlay(route_role_color), + overlay(face_color), + overlay(layer_face_color), + ] + if surface_color is not None: + overlay_images.append(overlay(surface_color)) + overlay_preview = torch.cat(overlay_images, dim=0) + overlay_output.parent.mkdir(parents=True, exist_ok=True) + save_image(overlay_preview.clamp(0.0, 1.0).detach().cpu(), overlay_output, nrow=view_count) + + +def build_arg_parser(): + parser = argparse.ArgumentParser(description="Infer UV conditioning with a dense UV parser.") + parser.add_argument("--parser_checkpoint", required=True) + parser.add_argument( + "--foreground_method", + choices=["flood", "legacy"], + default="flood", + help=( + "Background removal before dense parsing. 'flood' uses the top-left " + "pixel as a connected-color seed; 'legacy' leaves removal to the " + "former routing fallback." + ), + ) + parser.add_argument( + "--foreground_flood_tolerance", + type=float, + default=0.03, + help="Maximum per-channel RGB distance from the top-left flood seed.", + ) + parser.add_argument( + "--foreground_parser_background", + choices=["adaptive", "neutral"], + default="adaptive", + help="Solid background used for the masked RGB passed to dense parser.", + ) + parser.add_argument( + "--foreground_probability_output", + default="outputs/foreground_probability.png", + help=( + "Grayscale foreground score produced before dense parsing; this is " + "a binary mask when top-left flood fill is selected." + ), + ) + parser.add_argument( + "--foreground_mask_output", + default="outputs/foreground_mask.png", + help="Thresholded fixed-view foreground mask used by dense parsing.", + ) + parser.add_argument( + "--foreground_raw_mask_output", + default="outputs/foreground_mask_raw.png", + help="Binary foreground mask produced directly by top-left flood fill.", + ) + parser.add_argument( + "--foreground_cutout_output", + default="outputs/foreground_cutout.png", + help="Input views with the predicted background removed.", + ) + parser.add_argument( + "--foreground_parser_input_output", + default="outputs/foreground_parser_input.png", + help="Exact adaptive-background RGB images passed to dense parser.", + ) + parser.add_argument( + "--output", + default="outputs/pred_uv.png", + help=( + "Final deterministic RGBA UV PNG. Unknown inner-layer texels are " + "completed by the topology-aware simple inpainting algorithm; " + "the predicted outer layer is preserved unchanged." + ), + ) + parser.add_argument("--conditioning_output", default=None, help="Optional preview image for parser-splatted conditioning.") + parser.add_argument( + "--parser_uv_output", + default="outputs/parser_pred_uv.png", + help="Optional preliminary RGBA skin merged directly from parser conditioning.", + ) + parser.add_argument( + "--simple_inpaint_output", + default="outputs/parser_pred_uv_simple_inpainting.png", + help=( + "Optional copy of the deterministic pre-final UV repair. This is " + "the same completion used by --output." + ), + ) + parser.add_argument( + "--simple_inpaint_render_output", + default="outputs/simple_inpaint_render.png", + help=( + "Combined front+back render of the simple-inpainted skin. " + "Left half = front_left view, right half = back_left view." + ), + ) + parser.add_argument("--debug_output", default=None, help="Optional path to write a debug preview grid of predictions.") + parser.add_argument("--overlay_output", default=None, help="Optional path for segmentation overlays on canonicalized input views.") + parser.add_argument("--overlay_alpha", type=float, default=0.45) + parser.add_argument("--inner_cutout_output", default=None, help="Original-color cutout for routed inner-layer pixels.") + parser.add_argument("--outer_cutout_output", default=None, help="Original-color cutout for routed outer/decor pixels.") + parser.add_argument( + "--secondary_cutout_output", + default=None, + help="Original-color cutout for secondary/deeper surface pixels.", + ) + parser.add_argument( + "--color_source_output", + default=None, + help="Original RGB pixels permitted to contribute colors to parser UV output.", + ) + parser.add_argument("--face_output", default=None, help="Six-class routed cube-face visualization.") + parser.add_argument("--layer_face_output", default=None, help="Twelve-class inner/outer-by-face visualization.") + parser.add_argument("--raw_face_output", default=None, help="Six-class raw face-head visualization.") + parser.add_argument("--raw_layer_face_output", default=None, help="Twelve-class raw joint-head visualization.") + parser.add_argument( + "--canonical_foreground_output", + default=None, + help=( + "Foreground mask after coverage-preserving affine " + "canonicalization, before route filtering." + ), + ) + parser.add_argument("--geometry_grid_output", default=None, help="Fitted inner/outer cuboid UV grid preview.") + parser.add_argument( + "--geometry_overlay_output", + default=None, + help="Inner/outer fitted UV texel grids overlaid on canonicalized source views.", + ) + parser.add_argument( + "--geometry_routed_overlay_output", + default=None, + help="Inner/outer UV grids overlaid on only the pixels routed to that layer.", + ) + parser.add_argument("--geometry_fill_output", default=None, help="Classified RGB filled onto inner/outer cuboid grids.") + parser.add_argument( + "--outer_uv_occupancy_output", + default=None, + help=( + "Projected outer-UV occupancy diagnostic. Columns show raw " + "probability, topology-propagated probability, and component " + "routing masks (red=rejected, green=accepted, blue=seed)." + ), + ) + parser.add_argument("--front", default=None) + parser.add_argument("--back", default=None) + parser.add_argument("--combined", default=None) + parser.add_argument("--view_images", nargs="*", default=None) + parser.add_argument("--mappings_dir", default=None) + parser.add_argument("--fg_threshold", type=float, default=0.5) + parser.add_argument( + "--background_color_tolerance", + type=float, + default=48.0 / 255.0, + help="RGB distance used to reject solid-background and antialiased edge pixels.", + ) + parser.add_argument( + "--color_background_tolerance", + type=float, + default=8.0 / 255.0, + help=( + "Reject background-like RGB candidates only on the foreground " + "boundary before inverse UV color sampling." + ), + ) + parser.add_argument( + "--color_foreground_inset", + type=int, + default=1, + help="Foreground boundary width demoted during texel-center color selection.", + ) + parser.add_argument("--route_confidence_threshold", type=float, default=0.0) + parser.add_argument("--route_margin_threshold", type=float, default=0.0) + parser.add_argument("--outer_route_confidence_threshold", type=float, default=0.55) + parser.add_argument("--outer_route_margin_threshold", type=float, default=0.35) + parser.add_argument( + "--outer_uv_min_coverage", + type=float, + default=None, + help="Reject outer UV texels supported by less than this fraction of their projected footprint.", + ) + parser.add_argument( + "--outer_uv_min_source_pixels", + type=int, + default=33, + help="Minimum routed source pixels required to keep an outer UV texel.", + ) + parser.add_argument( + "--outer_silhouette_consistency", + dest="outer_silhouette_consistency", + action="store_true", + default=True, + help=( + "Reject an outer texel when its protruding projection is absent " + "from the observed foreground silhouette." + ), + ) + parser.add_argument( + "--no_outer_silhouette_consistency", + dest="outer_silhouette_consistency", + action="store_false", + ) + parser.add_argument( + "--outer_silhouette_min_coverage", type=float, default=0.50 + ) + parser.add_argument( + "--outer_silhouette_dilation", type=int, default=0 + ) + parser.add_argument( + "--outer_silhouette_min_pixels", type=int, default=4 + ) + parser.add_argument( + "--outer_geometry_rescue", + dest="outer_geometry_rescue", + action="store_true", + default=True, + help="Relax outer gates only for UV texels proven by outer-only silhouette or an exact secondary slot.", + ) + parser.add_argument( + "--no_outer_geometry_rescue", + dest="outer_geometry_rescue", + action="store_false", + ) + parser.add_argument( + "--outer_semantic_rescue", + dest="outer_semantic_rescue", + action="store_true", + default=True, + help="Relax outer gates only on parts whose global semantic heads predict a substantial outer layer.", + ) + parser.add_argument( + "--no_outer_semantic_rescue", + dest="outer_semantic_rescue", + action="store_false", + ) + parser.add_argument("--outer_semantic_presence_threshold", type=float, default=0.80) + parser.add_argument("--outer_semantic_coverage_threshold", type=float, default=0.20) + parser.add_argument("--outer_rescue_confidence_threshold", type=float, default=0.60) + parser.add_argument("--outer_rescue_margin_threshold", type=float, default=0.25) + parser.add_argument("--outer_rescue_min_coverage", type=float, default=0.10) + parser.add_argument( + "--geometry_route_texel_consensus", + dest="geometry_route_texel_consensus", + action="store_true", + default=None, + help="Use projected UV-cell voting instead of semantic-first per-pixel routing.", + ) + parser.add_argument( + "--no_geometry_route_texel_consensus", + dest="geometry_route_texel_consensus", + action="store_false", + ) + parser.add_argument( + "--geometry_route_texel_consensus_weight", + type=float, + default=None, + help="Blend weight for center-weighted texel consensus; 0 keeps local probabilities and 1 uses only the cell aggregate.", + ) + parser.add_argument( + "--geometry_route_preserve_outer_confidence", + type=float, + default=None, + help=( + "Minimum local outer confidence required in addition to the fused " + "texel gate when the raw route predicts outer." + ), + ) + parser.add_argument( + "--geometry_route_preserve_outer_margin", + type=float, + default=None, + help=( + "Minimum local outer margin required in addition to the fused " + "texel gate when the raw route predicts outer." + ), + ) + parser.add_argument( + "--geometry_route_consensus_outer_confidence", + type=float, + default=None, + help="Minimum fused confidence required to promote a route to outer.", + ) + parser.add_argument( + "--geometry_route_consensus_outer_margin", + type=float, + default=None, + help="Minimum fused margin ratio required to promote a route to outer.", + ) + parser.add_argument( + "--geometry_cross_view_outer_consistency", + dest="geometry_cross_view_outer_consistency", + action="store_true", + default=None, + help=( + "Veto a strong outer route when another view gives strong " + "background/inner evidence for the same outer UV texel." + ), + ) + parser.add_argument( + "--no_geometry_cross_view_outer_consistency", + dest="geometry_cross_view_outer_consistency", + action="store_false", + ) + parser.add_argument( + "--geometry_cross_view_outer_weight", type=float, default=None + ) + parser.add_argument( + "--geometry_cross_view_outer_positive_confidence", + type=float, + default=None, + ) + parser.add_argument( + "--geometry_cross_view_outer_positive_margin", + type=float, + default=None, + ) + parser.add_argument( + "--geometry_cross_view_outer_negative_confidence", + type=float, + default=None, + ) + parser.add_argument( + "--geometry_cross_view_outer_negative_margin", + type=float, + default=None, + ) + parser.add_argument( + "--geometry_cross_view_outer_background_max_coverage", + type=float, + default=None, + ) + parser.add_argument( + "--geometry_cross_view_outer_min_views", type=int, default=None + ) + parser.add_argument( + "--outer_uv_occupancy", + dest="outer_uv_occupancy", + action="store_true", + default=None, + help="Use the checkpoint's grouped 64x64 outer-layer occupancy prior.", + ) + parser.add_argument( + "--no_outer_uv_occupancy", + dest="outer_uv_occupancy", + action="store_false", + ) + parser.add_argument( + "--outer_uv_occupancy_blend_weight", type=float, default=None + ) + parser.add_argument( + "--outer_uv_occupancy_gate_threshold", type=float, default=None + ) + parser.add_argument( + "--outer_uv_occupancy_rescue_threshold", type=float, default=None + ) + parser.add_argument( + "--outer_uv_occupancy_rescue_route_threshold", + type=float, + default=None, + ) + parser.add_argument( + "--outer_uv_occupancy_auto_reliability", + dest="outer_uv_occupancy_auto_reliability", + action="store_true", + default=True, + help=( + "Allow occupancy to alter routing only when checkpoint validation " + "precision and recall clear the configured floors." + ), + ) + parser.add_argument( + "--no_outer_uv_occupancy_auto_reliability", + dest="outer_uv_occupancy_auto_reliability", + action="store_false", + ) + parser.add_argument( + "--outer_uv_occupancy_min_precision", type=float, default=0.60 + ) + parser.add_argument( + "--outer_uv_occupancy_min_recall", type=float, default=0.40 + ) + parser.add_argument( + "--outer_uv_component_routing", + dest="outer_uv_component_routing", + action="store_true", + default=None, + ) + parser.add_argument( + "--no_outer_uv_component_routing", + dest="outer_uv_component_routing", + action="store_false", + ) + parser.add_argument( + "--outer_uv_component_seed_threshold", type=float, default=None + ) + parser.add_argument( + "--outer_uv_component_grow_threshold", type=float, default=None + ) + parser.add_argument( + "--outer_uv_component_min_size", type=int, default=None + ) + parser.add_argument( + "--color_aggregation", + choices=SPLAT_COLOR_AGGREGATIONS, + default="grid_mode", + help="How colors inside each fitted layer/UV grid cell are selected.", + ) + parser.add_argument( + "--allow_semantic_fallback", + action="store_true", + help="Keep pixels whose strict semantic routing had no valid candidate.", + ) + parser.add_argument("--no_semantic_gate", dest="semantic_gate", action="store_false", default=None) + parser.add_argument("--affine_refine", dest="affine_refine", action="store_true", default=None) + parser.add_argument("--no_affine_refine", dest="affine_refine", action="store_false") + parser.add_argument("--affine_refine_translation_px", type=float, default=None) + parser.add_argument("--affine_refine_scale", type=float, default=None) + parser.add_argument("--alpha_threshold", type=float, default=0.5) + parser.add_argument("--device", default="auto") + return parser + + +def main(): + args = build_arg_parser().parse_args() + if not 0.0 <= args.background_color_tolerance <= 1.0: + raise ValueError("--background_color_tolerance must be in [0, 1].") + if not 0.0 <= args.color_background_tolerance <= 1.0: + raise ValueError("--color_background_tolerance must be in [0, 1].") + if args.color_foreground_inset < 0: + raise ValueError("--color_foreground_inset must be non-negative.") + if args.outer_uv_min_source_pixels < 1: + raise ValueError("--outer_uv_min_source_pixels must be positive.") + if not 0.0 <= args.outer_silhouette_min_coverage <= 1.0: + raise ValueError( + "--outer_silhouette_min_coverage must be in [0, 1]." + ) + if args.outer_silhouette_dilation < 0: + raise ValueError("--outer_silhouette_dilation must be non-negative.") + if args.outer_silhouette_min_pixels < 1: + raise ValueError("--outer_silhouette_min_pixels must be positive.") + if not 0.0 <= args.foreground_flood_tolerance <= 1.0: + raise ValueError("--foreground_flood_tolerance must be in [0, 1].") + if not 0.0 <= args.outer_uv_occupancy_min_precision <= 1.0: + raise ValueError( + "--outer_uv_occupancy_min_precision must be in [0, 1]." + ) + if not 0.0 <= args.outer_uv_occupancy_min_recall <= 1.0: + raise ValueError( + "--outer_uv_occupancy_min_recall must be in [0, 1]." + ) + if not any( + ( + args.output, + args.conditioning_output, + args.parser_uv_output, + args.simple_inpaint_output, + args.debug_output, + args.overlay_output, + args.inner_cutout_output, + args.outer_cutout_output, + args.secondary_cutout_output, + args.color_source_output, + args.face_output, + args.layer_face_output, + args.raw_face_output, + args.raw_layer_face_output, + args.geometry_grid_output, + args.geometry_overlay_output, + args.geometry_routed_overlay_output, + args.geometry_fill_output, + args.outer_uv_occupancy_output, + args.simple_inpaint_render_output, + ) + ): + raise ValueError( + "Provide --output, --conditioning_output, --parser_uv_output, --simple_inpaint_output, " + "--simple_inpaint_render_output, " + "--debug_output, --overlay_output, " + "--inner_cutout_output, --outer_cutout_output, --secondary_cutout_output, " + "and/or --color_source_output." + ) + + device = get_device(args.device) + parser_model, parser_args = load_parser(args.parser_checkpoint, device) + views = parse_views(parser_args.get("views", "walk_front_both_layer_ortho,walk_back_both_layer_ortho")) + if parser_model.view_classes not in (0, len(views)): + raise ValueError( + f"Parser checkpoint expects {parser_model.view_classes} views, but its metadata lists {len(views)}: {views}" + ) + mappings_dir = args.mappings_dir or parser_args.get("mappings_dir") + renderer = DifferentiableRenderer(mappings_dir=mappings_dir).to(device) + missing_views = [view for view in views if view not in renderer.views] + if missing_views: + raise ValueError(f"Unknown renderer views {missing_views}. Available views: {', '.join(renderer.views)}") + if parser_model.predict_affine and parser_model.surface_classes > 0: + mapping_surface_classes = surface_class_count(renderer, views) + if parser_model.surface_classes != mapping_surface_classes: + raise ValueError( + "Parser/mapping surface-slot mismatch: " + f"checkpoint={parser_model.surface_classes}, mappings={mapping_surface_classes}." + ) + + bg_color = parser_args.get("bg_color", (128, 128, 128)) + semantic_gate = parser_args.get("semantic_gate", True) if args.semantic_gate is None else args.semantic_gate + affine_refine = parser_args.get("affine_refine", True) if args.affine_refine is None else args.affine_refine + checkpoint_translation_px = parser_args.get("affine_refine_translation_px") + affine_refine_translation_px = ( + args.affine_refine_translation_px + if args.affine_refine_translation_px is not None + else 8.0 if checkpoint_translation_px is None else checkpoint_translation_px + ) + checkpoint_scale = parser_args.get("affine_refine_scale") + affine_refine_scale = ( + args.affine_refine_scale + if args.affine_refine_scale is not None + else 0.0 if checkpoint_scale is None else checkpoint_scale + ) + geometry_route_texel_consensus = ( + parser_args.get("geometry_route_texel_consensus", False) + if args.geometry_route_texel_consensus is None + else args.geometry_route_texel_consensus + ) + geometry_route_texel_consensus_weight = ( + parser_args.get("geometry_route_texel_consensus_weight", 0.60) + if args.geometry_route_texel_consensus_weight is None + else args.geometry_route_texel_consensus_weight + ) + geometry_route_preserve_outer_confidence = ( + parser_args.get("geometry_route_preserve_outer_confidence", 0.80) + if args.geometry_route_preserve_outer_confidence is None + else args.geometry_route_preserve_outer_confidence + ) + geometry_route_preserve_outer_margin = ( + parser_args.get("geometry_route_preserve_outer_margin", 0.35) + if args.geometry_route_preserve_outer_margin is None + else args.geometry_route_preserve_outer_margin + ) + geometry_route_consensus_outer_confidence = ( + parser_args.get("geometry_route_consensus_outer_confidence", 0.70) + if args.geometry_route_consensus_outer_confidence is None + else args.geometry_route_consensus_outer_confidence + ) + geometry_route_consensus_outer_margin = ( + parser_args.get("geometry_route_consensus_outer_margin", 0.20) + if args.geometry_route_consensus_outer_margin is None + else args.geometry_route_consensus_outer_margin + ) + geometry_cross_view_outer_consistency = ( + parser_args.get("geometry_cross_view_outer_consistency", False) + if args.geometry_cross_view_outer_consistency is None + else args.geometry_cross_view_outer_consistency + ) + geometry_cross_view_outer_weight = ( + parser_args.get("geometry_cross_view_outer_weight", 0.50) + if args.geometry_cross_view_outer_weight is None + else args.geometry_cross_view_outer_weight + ) + geometry_cross_view_outer_positive_confidence = ( + parser_args.get( + "geometry_cross_view_outer_positive_confidence", 0.70 + ) + if args.geometry_cross_view_outer_positive_confidence is None + else args.geometry_cross_view_outer_positive_confidence + ) + geometry_cross_view_outer_positive_margin = ( + parser_args.get("geometry_cross_view_outer_positive_margin", 0.20) + if args.geometry_cross_view_outer_positive_margin is None + else args.geometry_cross_view_outer_positive_margin + ) + geometry_cross_view_outer_negative_confidence = ( + parser_args.get( + "geometry_cross_view_outer_negative_confidence", 0.70 + ) + if args.geometry_cross_view_outer_negative_confidence is None + else args.geometry_cross_view_outer_negative_confidence + ) + geometry_cross_view_outer_negative_margin = ( + parser_args.get("geometry_cross_view_outer_negative_margin", 0.20) + if args.geometry_cross_view_outer_negative_margin is None + else args.geometry_cross_view_outer_negative_margin + ) + geometry_cross_view_outer_background_max_coverage = ( + parser_args.get( + "geometry_cross_view_outer_background_max_coverage", 0.25 + ) + if args.geometry_cross_view_outer_background_max_coverage is None + else args.geometry_cross_view_outer_background_max_coverage + ) + geometry_cross_view_outer_min_views = ( + parser_args.get("geometry_cross_view_outer_min_views", 2) + if args.geometry_cross_view_outer_min_views is None + else args.geometry_cross_view_outer_min_views + ) + outer_uv_occupancy_requested = ( + parser_args.get( + "outer_uv_occupancy_routing", + False, + ) + if args.outer_uv_occupancy is None + else args.outer_uv_occupancy + ) + occupancy_checkpoint_precision = parser_args.get( + "_checkpoint_outer_uv_occupancy_precision" + ) + occupancy_checkpoint_recall = parser_args.get( + "_checkpoint_outer_uv_occupancy_recall" + ) + occupancy_reliable = ( + occupancy_checkpoint_precision is not None + and occupancy_checkpoint_recall is not None + and occupancy_checkpoint_precision + >= float(args.outer_uv_occupancy_min_precision) + and occupancy_checkpoint_recall + >= float(args.outer_uv_occupancy_min_recall) + ) + outer_uv_occupancy = bool(outer_uv_occupancy_requested) and ( + occupancy_reliable + or not args.outer_uv_occupancy_auto_reliability + ) + print( + "outer_uv_occupancy_reliability=" + + json.dumps( + { + "auto": bool(args.outer_uv_occupancy_auto_reliability), + "enabled": bool(outer_uv_occupancy), + "requested": bool(outer_uv_occupancy_requested), + "checkpoint_precision": occupancy_checkpoint_precision, + "checkpoint_recall": occupancy_checkpoint_recall, + "min_precision": float( + args.outer_uv_occupancy_min_precision + ), + "min_recall": float(args.outer_uv_occupancy_min_recall), + }, + sort_keys=True, + ) + ) + outer_uv_occupancy_blend_weight = ( + parser_args.get("outer_uv_occupancy_blend_weight", 0.0) + if args.outer_uv_occupancy_blend_weight is None + else args.outer_uv_occupancy_blend_weight + ) + outer_uv_occupancy_gate_threshold = ( + parser_args.get("outer_uv_occupancy_gate_threshold", 0.15) + if args.outer_uv_occupancy_gate_threshold is None + else args.outer_uv_occupancy_gate_threshold + ) + outer_uv_occupancy_rescue_threshold = ( + parser_args.get("outer_uv_occupancy_rescue_threshold", 0.70) + if args.outer_uv_occupancy_rescue_threshold is None + else args.outer_uv_occupancy_rescue_threshold + ) + outer_uv_occupancy_rescue_route_threshold = ( + parser_args.get( + "outer_uv_occupancy_rescue_route_threshold", 0.30 + ) + if args.outer_uv_occupancy_rescue_route_threshold is None + else args.outer_uv_occupancy_rescue_route_threshold + ) + outer_uv_component_routing = ( + parser_args.get("outer_uv_component_routing", False) + if args.outer_uv_component_routing is None + else args.outer_uv_component_routing + ) + outer_uv_component_seed_threshold = ( + parser_args.get("outer_uv_component_seed_threshold", 0.80) + if args.outer_uv_component_seed_threshold is None + else args.outer_uv_component_seed_threshold + ) + outer_uv_component_grow_threshold = ( + parser_args.get("outer_uv_component_grow_threshold", 0.50) + if args.outer_uv_component_grow_threshold is None + else args.outer_uv_component_grow_threshold + ) + outer_uv_component_min_size = ( + parser_args.get("outer_uv_component_min_size", 2) + if args.outer_uv_component_min_size is None + else args.outer_uv_component_min_size + ) + outer_uv_min_coverage = ( + parser_args.get("outer_uv_min_coverage", 0.0) + if args.outer_uv_min_coverage is None + else args.outer_uv_min_coverage + ) + rendered = load_view_images(args, views, renderer, bg_color=bg_color).to(device) + view_ids = torch.arange(len(views), device=device) + with torch.no_grad(): + observed_foreground = None + parser_rendered = rendered + if args.foreground_method == "flood": + observed_foreground = estimate_top_left_flood_foreground( + rendered, + color_tolerance=args.foreground_flood_tolerance, + ) + foreground_log = { + "method": "top_left_flood", + "seed_rgb": [ + [int(round(channel * 255.0)) for channel in color] + for color in rendered[:, :3, 0, 0].detach().cpu().tolist() + ], + "tolerance": round(float(args.foreground_flood_tolerance), 6), + } + observed_foreground = save_flood_outputs( + rendered, + observed_foreground, + view_count=len(views), + probability_output=args.foreground_probability_output, + raw_mask_output=args.foreground_raw_mask_output, + mask_output=args.foreground_mask_output, + cutout_output=args.foreground_cutout_output, + ) + ( + parser_rendered, + parser_background_rgb, + parser_background_indices, + ) = build_parser_input( + rendered, + observed_foreground, + bg_color=bg_color, + background_mode=args.foreground_parser_background, + return_background=True, + ) + if args.foreground_parser_input_output: + parser_input_path = Path(args.foreground_parser_input_output) + parser_input_path.parent.mkdir(parents=True, exist_ok=True) + save_image( + parser_rendered[:, :3].detach().cpu(), + parser_input_path, + nrow=len(views), + ) + foreground_log.update( + { + "kept_pixels": int(observed_foreground.sum().item()), + "rejected_background_pixels": int( + (~observed_foreground).sum().item() + ), + "parser_background_mode": args.foreground_parser_background, + "parser_background_rgb": [ + [int(round(channel * 255.0)) for channel in color] + for color in parser_background_rgb.detach().cpu().tolist() + ], + "parser_background_indices": parser_background_indices, + } + ) + print( + "foreground_filter=" + + json.dumps(foreground_log, sort_keys=True) + ) + # Background removal affects parser features, while routing/splatting + # below still reads exact foreground colors from the original input. + outputs = parser_model( + parser_rendered, + view_ids=view_ids, + semantic_foreground=observed_foreground, + ) + outputs = attach_projected_outer_uv_occupancy( + parser_model, + outputs, + renderer, + views, + observed_foreground=observed_foreground, + center_power=float( + parser_args.get("route_texel_center_power", 2.0) + ), + ) + if ( + args.outer_uv_occupancy_output + and "outer_uv_occupancy_logits" in outputs + ): + occupancy_probability = torch.sigmoid( + outputs["outer_uv_occupancy_logits"] + ) + occupancy_components = outer_uv_topology_hysteresis( + occupancy_probability, + seed_threshold=outer_uv_component_seed_threshold, + grow_threshold=outer_uv_component_grow_threshold, + min_component_size=outer_uv_component_min_size, + ) + raw_preview = occupancy_probability.expand(-1, 3, -1, -1) + structured_preview = occupancy_components[ + "probability" + ].expand(-1, 3, -1, -1) + component_preview = torch.cat( + ( + occupancy_components["rejected_candidate"].float(), + occupancy_components["accepted"].float(), + occupancy_components["seed"].float(), + ), + dim=1, + ) + occupancy_preview = torch.cat( + (raw_preview, structured_preview, component_preview), + dim=-1, + ) + occupancy_path = Path(args.outer_uv_occupancy_output) + occupancy_path.parent.mkdir(parents=True, exist_ok=True) + save_image(occupancy_preview.detach().cpu(), occupancy_path) + print(f"Saved outer_uv_occupancy={occupancy_path}") + conditioning, routing_details = splat_parser_predictions_to_uv_conditioning( + rendered, + outputs, + renderer=renderer, + views=views, + group_size=len(views), + fg_threshold=args.fg_threshold, + bg_color=bg_color, + semantic_gate=semantic_gate, + affine_refine=affine_refine, + affine_refine_translation_px=affine_refine_translation_px, + affine_refine_scale=affine_refine_scale, + route_confidence_threshold=args.route_confidence_threshold, + route_margin_threshold=args.route_margin_threshold, + outer_route_confidence_threshold=args.outer_route_confidence_threshold, + outer_route_margin_threshold=args.outer_route_margin_threshold, + outer_uv_min_coverage=outer_uv_min_coverage, + outer_uv_min_source_pixels=args.outer_uv_min_source_pixels, + outer_silhouette_consistency=args.outer_silhouette_consistency, + outer_silhouette_min_coverage=( + args.outer_silhouette_min_coverage + ), + outer_silhouette_dilation=args.outer_silhouette_dilation, + outer_silhouette_min_pixels=args.outer_silhouette_min_pixels, + outer_geometry_rescue=args.outer_geometry_rescue, + outer_semantic_rescue=args.outer_semantic_rescue, + outer_semantic_presence_threshold=args.outer_semantic_presence_threshold, + outer_semantic_coverage_threshold=args.outer_semantic_coverage_threshold, + outer_rescue_confidence_threshold=args.outer_rescue_confidence_threshold, + outer_rescue_margin_threshold=args.outer_rescue_margin_threshold, + outer_rescue_min_coverage=args.outer_rescue_min_coverage, + color_aggregation=args.color_aggregation, + geometry_route_texel_consensus=geometry_route_texel_consensus, + geometry_route_texel_consensus_weight=( + geometry_route_texel_consensus_weight + ), + geometry_route_preserve_outer_confidence=( + geometry_route_preserve_outer_confidence + ), + geometry_route_preserve_outer_margin=( + geometry_route_preserve_outer_margin + ), + geometry_route_consensus_outer_confidence=( + geometry_route_consensus_outer_confidence + ), + geometry_route_consensus_outer_margin=( + geometry_route_consensus_outer_margin + ), + geometry_cross_view_outer_consistency=( + geometry_cross_view_outer_consistency + ), + geometry_cross_view_outer_weight=( + geometry_cross_view_outer_weight + ), + geometry_cross_view_outer_positive_confidence=( + geometry_cross_view_outer_positive_confidence + ), + geometry_cross_view_outer_positive_margin=( + geometry_cross_view_outer_positive_margin + ), + geometry_cross_view_outer_negative_confidence=( + geometry_cross_view_outer_negative_confidence + ), + geometry_cross_view_outer_negative_margin=( + geometry_cross_view_outer_negative_margin + ), + geometry_cross_view_outer_background_max_coverage=( + geometry_cross_view_outer_background_max_coverage + ), + geometry_cross_view_outer_min_views=( + geometry_cross_view_outer_min_views + ), + outer_uv_occupancy=outer_uv_occupancy, + outer_uv_occupancy_blend_weight=( + outer_uv_occupancy_blend_weight + ), + outer_uv_occupancy_gate_threshold=( + outer_uv_occupancy_gate_threshold + ), + outer_uv_occupancy_rescue_threshold=( + outer_uv_occupancy_rescue_threshold + ), + outer_uv_occupancy_rescue_route_threshold=( + outer_uv_occupancy_rescue_route_threshold + ), + outer_uv_component_routing=outer_uv_component_routing, + outer_uv_component_seed_threshold=( + outer_uv_component_seed_threshold + ), + outer_uv_component_grow_threshold=( + outer_uv_component_grow_threshold + ), + outer_uv_component_min_size=outer_uv_component_min_size, + observed_foreground=observed_foreground, + background_color_tolerance=args.background_color_tolerance, + color_background_tolerance=args.color_background_tolerance, + color_foreground_inset=args.color_foreground_inset, + reject_semantic_fallback=not args.allow_semantic_fallback, + include_rejected_context=False, + include_confidence=False, + return_details=True, + ) + + routing = routing_details.get("routing") + if routing is not None: + if args.canonical_foreground_output: + canonical_foreground_path = Path( + args.canonical_foreground_output + ) + canonical_foreground_path.parent.mkdir( + parents=True, exist_ok=True + ) + save_image( + routing["observed_foreground"] + .unsqueeze(1) + .to(dtype=rendered.dtype) + .detach() + .cpu(), + canonical_foreground_path, + nrow=len(views), + ) + observed_routing_foreground = routing.get( + "observed_foreground", routing["raw_foreground"] + ) + observed_count = int(observed_routing_foreground.sum().item()) + canonical_foreground_coverage_rescued_count = int( + routing.get( + "canonical_foreground_coverage_rescued", + torch.zeros_like(observed_routing_foreground), + ).sum().item() + ) + unrouted_observed_count = int( + ( + observed_routing_foreground + & ~routing["raw_foreground"] + ).sum().item() + ) + raw_count = int(routing["raw_foreground"].sum().item()) + rejected_count = int(routing["rejected"].sum().item()) + kept = routing["foreground"] + kept_inner_count = int((kept & (routing["layer"] == 0)).sum().item()) + kept_outer_count = int((kept & (routing["layer"] == 1)).sum().item()) + raw_inner = routing["raw_foreground"] & (routing["layer"] == 0) + raw_outer = routing["raw_foreground"] & (routing["layer"] == 1) + rejected_inner = routing["rejected"] & (routing["layer"] == 0) + rejected_outer = routing["rejected"] & (routing["layer"] == 1) + outer_confidence = routing["confidence"][raw_outer].float() + outer_margin_ratio = routing["confidence_margin_ratio"][raw_outer].float() + + def quantile_or_zero(values, quantile): + return ( + float(torch.quantile(values, quantile).item()) + if values.numel() > 0 + else 0.0 + ) + coverage_rejected_outer = raw_outer & ( + routing.get("outer_uv_coverage", torch.ones_like(routing["confidence"])) + < routing.get( + "outer_required_coverage", + torch.full_like(routing["confidence"], outer_uv_min_coverage), + ) + ) + outer_silhouette_assessed_count = int( + routing.get( + "outer_silhouette_assessed", + torch.zeros_like(raw_outer), + ).sum().item() + ) + outer_silhouette_rejected_count = int( + routing.get( + "outer_silhouette_rejected", + torch.zeros_like(raw_outer), + ).sum().item() + ) + secondary_count = int(routing.get("secondary", torch.zeros_like(raw_outer)).sum().item()) + routed_secondary_count = int( + routing.get("secondary_routed", torch.zeros_like(raw_outer)).sum().item() + ) + rejected_secondary_count = int( + routing.get("secondary_rejected", torch.zeros_like(raw_outer)).sum().item() + ) + background_rejected_count = int( + routing.get("background_rejected", torch.zeros_like(raw_outer)).sum().item() + ) + color_rejected_count = int( + routing.get("color_rejected", torch.zeros_like(raw_outer)).sum().item() + ) + outer_source_rejected_count = int( + routing.get( + "outer_source_rejected", torch.zeros_like(raw_outer) + ).sum().item() + ) + outer_geometry_supported_count = int( + routing.get("outer_geometry_supported", torch.zeros_like(raw_outer)).sum().item() + ) + outer_geometry_rescued_count = int( + routing.get("outer_geometry_rescued", torch.zeros_like(raw_outer)).sum().item() + ) + outer_semantic_supported_count = int( + routing.get("outer_semantic_supported", torch.zeros_like(raw_outer)).sum().item() + ) + outer_semantic_rescued_count = int( + routing.get("outer_semantic_rescued", torch.zeros_like(raw_outer)).sum().item() + ) + rejected_context_count = int( + routing.get("rejected_context", torch.zeros_like(raw_outer)).sum().item() + ) + rejected_context_alpha_supported_count = int( + routing.get( + "rejected_context_alpha_supported", + torch.zeros_like(raw_outer), + ).sum().item() + ) + outer_geometry_supported_rejected_count = int( + ( + routing.get("outer_geometry_supported", torch.zeros_like(raw_outer)) + & routing["rejected"] + ).sum().item() + ) + raw_secondary_count = int( + ( + (routing.get("raw_route_role", routing.get("route_role")) == 2) + & (torch.sigmoid(routing_details["outputs"]["foreground"])[:, 0] > args.fg_threshold) + ).sum().item() + ) + consensus_changed_count = int( + routing.get( + "consensus_changed", torch.zeros_like(raw_outer) + ).sum().item() + ) + consensus_inner_to_outer_count = int( + routing.get( + "consensus_inner_to_outer", torch.zeros_like(raw_outer) + ).sum().item() + ) + consensus_outer_to_inner_count = int( + routing.get( + "consensus_outer_to_inner", torch.zeros_like(raw_outer) + ).sum().item() + ) + consensus_preserved_outer_count = int( + routing.get( + "consensus_preserved_outer", torch.zeros_like(raw_outer) + ).sum().item() + ) + consensus_outer_gate_rejected_count = int( + routing.get( + "consensus_outer_gate_rejected", + torch.zeros_like(raw_outer), + ).sum().item() + ) + consensus_outer_gate_deferred_count = int( + ( + routing.get( + "consensus_outer_gate_rejected", + torch.zeros_like(raw_outer), + ) + & (routing.get("route_role") == 1) + ).sum().item() + ) + cross_view_outer_shared_count = int( + routing.get( + "cross_view_outer_shared", torch.zeros_like(raw_outer) + ).sum().item() + ) + cross_view_outer_conflict_count = int( + routing.get( + "cross_view_outer_conflict", torch.zeros_like(raw_outer) + ).sum().item() + ) + cross_view_outer_vetoed_count = int( + routing.get( + "cross_view_outer_vetoed", torch.zeros_like(raw_outer) + ).sum().item() + ) + occupancy_promoted_outer_count = int( + routing.get( + "occupancy_rescued_outer", torch.zeros_like(raw_outer) + ).sum().item() + ) + occupancy_rejected_outer_count = int( + routing.get( + "occupancy_rejected_outer", torch.zeros_like(raw_outer) + ).sum().item() + ) + occupancy_supported_outer_count = int( + routing.get( + "outer_occupancy_supported", + torch.zeros_like(raw_outer), + ).sum().item() + ) + occupancy_trusted_outer_count = int( + routing.get( + "outer_occupancy_rescued", + torch.zeros_like(raw_outer), + ).sum().item() + ) + occupancy_component_seed_count = int( + routing.get( + "projected_outer_component_seed", + torch.zeros_like(raw_outer), + ).sum().item() + ) + occupancy_component_grown_count = int( + routing.get( + "projected_outer_component_grown", + torch.zeros_like(raw_outer), + ).sum().item() + ) + occupancy_component_rejected_count = int( + routing.get( + "projected_outer_component_rejected", + torch.zeros_like(raw_outer), + ).sum().item() + ) + print( + "routing_filter=" + + json.dumps( + { + "observed_foreground_pixels": observed_count, + "canonical_foreground_coverage_rescued_pixels": ( + canonical_foreground_coverage_rescued_count + ), + "unrouted_observed_pixels": unrouted_observed_count, + "raw_pixels": raw_count, + "kept_pixels": raw_count - rejected_count, + "kept_inner_pixels": kept_inner_count, + "kept_outer_pixels": kept_outer_count, + "kept_outer_percent": round( + 100.0 * kept_outer_count / max(kept_inner_count + kept_outer_count, 1), + 3, + ), + "rejected_pixels": rejected_count, + "rejected_percent": round(100.0 * rejected_count / max(raw_count, 1), 3), + "inner_rejected_percent": round( + 100.0 * int(rejected_inner.sum().item()) / max(int(raw_inner.sum().item()), 1), + 3, + ), + "outer_rejected_percent": round( + 100.0 * int(rejected_outer.sum().item()) / max(int(raw_outer.sum().item()), 1), + 3, + ), + "outer_confidence_p50": round( + quantile_or_zero(outer_confidence, 0.50), 6 + ), + "outer_confidence_p75": round( + quantile_or_zero(outer_confidence, 0.75), 6 + ), + "outer_confidence_p90": round( + quantile_or_zero(outer_confidence, 0.90), 6 + ), + "outer_margin_p50": round( + quantile_or_zero(outer_margin_ratio, 0.50), 6 + ), + "outer_margin_p75": round( + quantile_or_zero(outer_margin_ratio, 0.75), 6 + ), + "outer_coverage_rejected_pixels": int(coverage_rejected_outer.sum().item()), + "outer_source_rejected_pixels": outer_source_rejected_count, + "outer_uv_min_source_pixels": int( + args.outer_uv_min_source_pixels + ), + "outer_silhouette_consistency": bool( + args.outer_silhouette_consistency + ), + "outer_silhouette_min_coverage": float( + args.outer_silhouette_min_coverage + ), + "outer_silhouette_view_reduction": "minimum", + "outer_silhouette_dilation": int( + args.outer_silhouette_dilation + ), + "outer_silhouette_evidence_inset": int( + args.color_foreground_inset + ), + "outer_silhouette_assessed_pixels": ( + outer_silhouette_assessed_count + ), + "outer_silhouette_rejected_pixels": ( + outer_silhouette_rejected_count + ), + "outer_geometry_supported_pixels": outer_geometry_supported_count, + "outer_geometry_rescued_pixels": outer_geometry_rescued_count, + "outer_geometry_supported_rejected_pixels": outer_geometry_supported_rejected_count, + "outer_semantic_supported_pixels": outer_semantic_supported_count, + "outer_semantic_rescued_pixels": outer_semantic_rescued_count, + "rejected_context_pixels": rejected_context_count, + "rejected_context_alpha_supported_pixels": ( + rejected_context_alpha_supported_count + ), + "background_rejected_pixels": background_rejected_count, + "background_color_tolerance": round( + float(args.background_color_tolerance), 6 + ), + "color_source_rejected_pixels": color_rejected_count, + "color_background_tolerance": round( + float(args.color_background_tolerance), 6 + ), + "color_foreground_inset": int(args.color_foreground_inset), + "secondary_backface_pixels": secondary_count, + "routed_secondary_pixels": routed_secondary_count, + "rejected_secondary_pixels": rejected_secondary_count, + "raw_secondary_backface_pixels": raw_secondary_count, + "consensus_changed_pixels": consensus_changed_count, + "consensus_inner_to_outer_pixels": ( + consensus_inner_to_outer_count + ), + "consensus_outer_to_inner_pixels": ( + consensus_outer_to_inner_count + ), + "consensus_preserved_outer_pixels": ( + consensus_preserved_outer_count + ), + "consensus_outer_gate_rejected_pixels": ( + consensus_outer_gate_rejected_count + ), + "consensus_outer_gate_deferred_pixels": ( + consensus_outer_gate_deferred_count + ), + "cross_view_outer_consistency": bool( + geometry_cross_view_outer_consistency + ), + "cross_view_outer_shared_pixels": ( + cross_view_outer_shared_count + ), + "cross_view_outer_conflict_pixels": ( + cross_view_outer_conflict_count + ), + "cross_view_outer_vetoed_pixels": ( + cross_view_outer_vetoed_count + ), + "cross_view_outer_weight": round( + float(geometry_cross_view_outer_weight), 6 + ), + "consensus_weight": round( + float(geometry_route_texel_consensus_weight), 6 + ), + "outer_uv_occupancy_available": bool( + routing.get( + "outer_uv_occupancy_available", + torch.zeros_like(raw_outer), + ).any().item() + ), + "outer_uv_occupancy_promoted_pixels": ( + occupancy_promoted_outer_count + ), + "outer_uv_occupancy_rejected_pixels": ( + occupancy_rejected_outer_count + ), + "outer_uv_occupancy_supported_pixels": ( + occupancy_supported_outer_count + ), + "outer_uv_occupancy_trusted_rescue_pixels": ( + occupancy_trusted_outer_count + ), + "outer_uv_occupancy_blend_weight": round( + float(outer_uv_occupancy_blend_weight), 6 + ), + "outer_uv_component_routing": bool( + outer_uv_component_routing + ), + "outer_uv_component_seed_pixels": ( + occupancy_component_seed_count + ), + "outer_uv_component_grown_pixels": ( + occupancy_component_grown_count + ), + "outer_uv_component_rejected_pixels": ( + occupancy_component_rejected_count + ), + "outer_uv_component_seed_threshold": round( + float(outer_uv_component_seed_threshold), 6 + ), + "outer_uv_component_grow_threshold": round( + float(outer_uv_component_grow_threshold), 6 + ), + }, + sort_keys=True, + ) + ) + + alignment = routing_details.get("alignment") + if alignment is not None: + translation = alignment["translation_px"].detach().cpu() + scale_percent = alignment["scale_percent"].detach().cpu() + score_before = alignment["score_before"].detach().cpu() + score_after = alignment["score_after"].detach().cpu() + accepted = alignment["accepted"].detach().cpu() + for index in range(translation.shape[0]): + print( + "affine_refinement=" + + json.dumps( + { + "view": views[index % len(views)], + "accepted": bool(accepted[index]), + "dx_px": round(float(translation[index, 0]), 3), + "dy_px": round(float(translation[index, 1]), 3), + "scale_percent": round(float(scale_percent[index]), 4), + "score_before": round(float(score_before[index]), 6), + "score_after": round(float(score_after[index]), 6), + }, + sort_keys=True, + ) + ) + + if any( + ( + args.debug_output, + args.overlay_output, + args.inner_cutout_output, + args.outer_cutout_output, + args.secondary_cutout_output, + args.color_source_output, + args.face_output, + args.layer_face_output, + args.raw_face_output, + args.raw_layer_face_output, + args.geometry_grid_output, + args.geometry_overlay_output, + args.geometry_routed_overlay_output, + args.geometry_fill_output, + args.outer_uv_occupancy_output, + ) + ): + save_debug_preview( + routing_details["rendered"], + routing_details["outputs"], + len(views), + Path(args.debug_output) if args.debug_output else None, + args.fg_threshold, + bg_color=bg_color, + routing=routing_details["routing"], + overlay_output=Path(args.overlay_output) if args.overlay_output else None, + overlay_alpha=args.overlay_alpha, + inner_cutout_output=Path(args.inner_cutout_output) if args.inner_cutout_output else None, + outer_cutout_output=Path(args.outer_cutout_output) if args.outer_cutout_output else None, + secondary_cutout_output=( + Path(args.secondary_cutout_output) if args.secondary_cutout_output else None + ), + color_source_output=( + Path(args.color_source_output) if args.color_source_output else None + ), + face_output=Path(args.face_output) if args.face_output else None, + layer_face_output=Path(args.layer_face_output) if args.layer_face_output else None, + raw_face_output=Path(args.raw_face_output) if args.raw_face_output else None, + raw_layer_face_output=Path(args.raw_layer_face_output) if args.raw_layer_face_output else None, + geometry_grid_output=Path(args.geometry_grid_output) if args.geometry_grid_output else None, + geometry_overlay_output=( + Path(args.geometry_overlay_output) if args.geometry_overlay_output else None + ), + geometry_routed_overlay_output=( + Path(args.geometry_routed_overlay_output) + if args.geometry_routed_overlay_output + else None + ), + geometry_fill_output=Path(args.geometry_fill_output) if args.geometry_fill_output else None, + renderer=renderer, + views=views, + raw_outputs=outputs, + raw_observed_foreground=observed_foreground, + ) + + if args.conditioning_output: + save_conditioning_preview(conditioning.detach().cpu(), Path(args.conditioning_output)) + + if args.parser_uv_output: + save_parser_uv( + conditioning.detach().cpu(), + Path(args.parser_uv_output), + alpha_threshold=args.alpha_threshold, + enforce_base_alpha=False, + ) + + needs_repair = bool( + args.simple_inpaint_output + or args.output + or args.simple_inpaint_render_output + ) + repair_outputs = [] + if args.simple_inpaint_output: + repair_outputs.append( + ("parser_simple_inpaint_uv", Path(args.simple_inpaint_output)) + ) + if args.output: + repair_outputs.append(("completed_uv", Path(args.output))) + if needs_repair: + repaired, stats = simple_inpaint_uv( + conditioning.detach().cpu(), + alpha_threshold=args.alpha_threshold, + ) + print("simple_inpaint_stats=" + json.dumps(stats, sort_keys=True)) + written_paths = set() + for output_label, output_path in repair_outputs: + resolved_path = output_path.resolve() + if resolved_path not in written_paths: + output_path.parent.mkdir(parents=True, exist_ok=True) + tensor_to_rgba_image(repaired).save(output_path) + written_paths.add(resolved_path) + print(f"Saved {output_label}={output_path}") + + if args.simple_inpaint_render_output: + repaired_skin = repaired.unsqueeze(0).to(device) + rendered_views = renderer.forward(repaired_skin) + expected_views = ("front_left", "back_left") + missing = [v for v in expected_views if v not in rendered_views] + if missing: + print( + "WARNING: simple_inpaint_render_output requires views " + f"{expected_views}, but renderer only has " + f"{list(rendered_views.keys())}. " + f"Missing views: {missing}. " + "Skipping combined render." + ) + else: + bg = ( + torch.tensor(bg_color, device=device, dtype=repaired_skin.dtype) + .view(1, 3, 1, 1) + / 255.0 + ) + front_rgba = rendered_views["front_left"] + back_rgba = rendered_views["back_left"] + front_rgb = front_rgba[:, :3] * front_rgba[:, 3:4] + bg * ( + 1.0 - front_rgba[:, 3:4] + ) + back_rgb = back_rgba[:, :3] * back_rgba[:, 3:4] + bg * ( + 1.0 - back_rgba[:, 3:4] + ) + combined = torch.cat([front_rgb, back_rgb], dim=3) + render_path = Path(args.simple_inpaint_render_output) + render_path.parent.mkdir(parents=True, exist_ok=True) + save_image( + combined[0].clamp(0.0, 1.0).cpu(), + render_path, + ) + print(f"Saved simple_inpaint_render={render_path}") + + +if __name__ == "__main__": + main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/losses.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..f29c70fc0e16357044532480d2cba381dabc6108 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/losses.py @@ -0,0 +1,683 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from SkingToolkit.dense_uv_parser.utils import IGNORE_INDEX, combine_layer_face + + +def _deterministic_cross_entropy( + logits, + target, + weight=None, + ignore_index=IGNORE_INDEX, +): + """Cross entropy without CUDA's non-deterministic NLLLoss2d kernel.""" + if logits.dim() < 2: + raise ValueError("Cross-entropy logits must include a class dimension.") + expected_target_shape = logits.shape[:1] + logits.shape[2:] + if target.shape != expected_target_shape: + raise ValueError( + f"Expected target shape {expected_target_shape}, got " + f"{tuple(target.shape)}." + ) + + logits = logits.float() + target = target.long() + valid = target != int(ignore_index) + if not valid.any(): + return logits.sum() * 0.0 + + class_count = logits.shape[1] + safe_target = target.clamp(0, class_count - 1) + log_probability = F.log_softmax(logits, dim=1) + per_element = -log_probability.gather( + 1, + safe_target.unsqueeze(1), + ).squeeze(1) + if weight is None: + denominator = valid.sum().to(dtype=per_element.dtype) + else: + weight = torch.as_tensor( + weight, + device=logits.device, + dtype=per_element.dtype, + ) + if weight.numel() != class_count: + raise ValueError( + f"Expected {class_count} class weights, got {weight.numel()}." + ) + element_weight = weight.gather(0, safe_target.reshape(-1)).reshape_as( + safe_target + ) + per_element = per_element * element_weight + denominator = element_weight[valid].sum() + return per_element[valid].sum() / denominator.clamp_min(1e-12) + + +def _balanced_cross_entropy( + logits, + target, + max_weight=4.0, + min_weight=0.0, + class_weight_caps=None, +): + """Give small projected faces useful gradient without letting them dominate.""" + if min_weight < 0 or max_weight < min_weight: + raise ValueError("Class weights require 0 <= min_weight <= max_weight.") + valid = target != IGNORE_INDEX + if not valid.any(): + return logits.new_tensor(0.0) + counts = torch.bincount(target[valid], minlength=logits.shape[1]).float() + active = counts > 0 + weights = torch.where(active, counts.clamp_min(1.0).rsqrt(), torch.zeros_like(counts)) + weights[active] /= weights[active].mean().clamp_min(1e-6) + weights[active] = weights[active].clamp(min=min_weight, max=max_weight) + if class_weight_caps is not None: + caps = torch.as_tensor( + class_weight_caps, + device=weights.device, + dtype=weights.dtype, + ) + if caps.numel() != logits.shape[1]: + raise ValueError( + f"Expected {logits.shape[1]} class-weight caps, got {caps.numel()}." + ) + if (caps < 0).any(): + raise ValueError("Class-weight caps must be non-negative.") + weights = torch.minimum(weights, caps) + return _deterministic_cross_entropy( + logits, + target, + weight=weights, + ignore_index=IGNORE_INDEX, + ) + + +def outer_false_positive_loss(logits, target, outer_index=1, gamma=2.0): + """Focal negative loss for non-outer pixels assigned outer probability. + + Ordinary class-balanced cross entropy can underweight abundant inner pixels + while increasing the rare outer target weight. This term restores the + asymmetric cost that matters for UV reconstruction: confidently inventing + an outer texel is worse than leaving an uncertain outer texel unknown. + """ + if gamma < 0: + raise ValueError("gamma must be non-negative.") + if not 0 <= outer_index < logits.shape[1]: + raise ValueError(f"outer_index={outer_index} is invalid for {logits.shape[1]} classes.") + negative = (target != IGNORE_INDEX) & (target != outer_index) + if not negative.any(): + return logits.new_zeros((), dtype=torch.float32) + outer_probability = torch.softmax(logits.float(), dim=1)[:, outer_index] + outer_probability = outer_probability.clamp(1e-6, 1.0 - 1e-6) + focal_negative = outer_probability.pow(gamma) * (-torch.log1p(-outer_probability)) + return focal_negative[negative].mean() + + +def outer_false_negative_loss(logits, target, outer_index=1, gamma=2.0): + """Focal positive loss for visible outer pixels assigned too little outer probability.""" + if gamma < 0: + raise ValueError("gamma must be non-negative.") + if not 0 <= outer_index < logits.shape[1]: + raise ValueError(f"outer_index={outer_index} is invalid for {logits.shape[1]} classes.") + positive = target == outer_index + if not positive.any(): + return logits.new_zeros((), dtype=torch.float32) + outer_probability = torch.softmax(logits.float(), dim=1)[:, outer_index] + outer_probability = outer_probability.clamp(1e-6, 1.0 - 1e-6) + focal_positive = (1.0 - outer_probability).pow(gamma) * (-outer_probability.log()) + return focal_positive[positive].mean() + + +def primary_route_swap_loss(logits, target, gamma=2.0): + """Macro-balanced focal loss for inner/outer swaps. + + Inner pixels are much more common than visible outer pixels. Computing one + mean per primary role before averaging prevents the abundant inner class + from deciding the gradient while still leaving secondary/backface handling + to the ordinary three-class route loss. + """ + if logits.shape[1] < 2: + raise ValueError("Primary route swap loss requires inner and outer logits.") + if gamma < 0: + raise ValueError("gamma must be non-negative.") + probability = torch.softmax(logits[:, :2].float(), dim=1).clamp( + 1e-6, 1.0 - 1e-6 + ) + class_losses = [] + for role in (0, 1): + selected = target == role + if not selected.any(): + continue + correct = probability[:, role][selected] + class_losses.append( + ((1.0 - correct).pow(gamma) * -correct.log()).mean() + ) + if not class_losses: + return logits.new_zeros((), dtype=torch.float32) + return torch.stack(class_losses).mean() + + +def projected_texel_consistency_loss(logits, target, uv, uv_size=64): + """Make pixels from the same projected UV texel agree on route role. + + Groups are kept per rendered item and ground-truth role. Including the + role in the key avoids forcing a rare legitimate secondary observation to + agree with a primary observation of the same atlas texel. + """ + if logits.dim() != 4 or target.shape != logits.shape[:1] + logits.shape[-2:]: + raise ValueError("Route logits and targets must be NCHW and NHW.") + if uv.shape != (logits.shape[0], 2, *logits.shape[-2:]): + raise ValueError( + f"Expected UV target shape {(logits.shape[0], 2, *logits.shape[-2:])}, " + f"got {tuple(uv.shape)}." + ) + if uv_size < 1: + raise ValueError("uv_size must be positive.") + valid = target != IGNORE_INDEX + if not valid.any(): + return logits.new_zeros((), dtype=torch.float32) + + probabilities = torch.softmax(logits.float(), dim=1) + x = (uv[:, 0].clamp(0.0, 1.0) * (uv_size - 1)).round().long() + y = (uv[:, 1].clamp(0.0, 1.0) * (uv_size - 1)).round().long() + flat_uv = y * uv_size + x + role_count = logits.shape[1] + uv_count = uv_size * uv_size + item = torch.arange(logits.shape[0], device=logits.device).view(-1, 1, 1) + safe_target = target.clamp(0, role_count - 1) + group = item * (role_count * uv_count) + safe_target * uv_count + flat_uv + + flat_probabilities = probabilities.permute(0, 2, 3, 1).reshape(-1, role_count) + flat_group = group.reshape(-1) + flat_valid = valid.reshape(-1) + selected_group = flat_group[flat_valid] + selected_probabilities = flat_probabilities[flat_valid] + group_count = logits.shape[0] * role_count * uv_count + sums = probabilities.new_zeros(group_count, role_count) + sums.index_add_(0, selected_group, selected_probabilities) + counts = probabilities.new_zeros(group_count, 1) + counts.index_add_( + 0, + selected_group, + probabilities.new_ones(selected_group.shape[0], 1), + ) + means = sums / counts.clamp_min(1.0) + selected_means = means.index_select(0, selected_group) + repeated = counts.index_select(0, selected_group)[:, 0] > 1 + if not repeated.any(): + return logits.new_zeros((), dtype=torch.float32) + return ( + selected_probabilities[repeated] - selected_means[repeated] + ).square().sum(dim=1).mean() + + +def route_prior_regularization(outputs, tv_weight=1.0): + """Keep the learned fixed-view prior smooth and weaker than image evidence.""" + prior = outputs.get("route_role_prior_raw") + if prior is None: + return outputs["layer"].new_zeros((), dtype=torch.float32) + prior = prior.float() + l2 = prior.square().mean() + tv_terms = [] + if prior.shape[-2] > 1: + tv_terms.append((prior[..., 1:, :] - prior[..., :-1, :]).square().mean()) + if prior.shape[-1] > 1: + tv_terms.append((prior[..., :, 1:] - prior[..., :, :-1]).square().mean()) + tv = torch.stack(tv_terms).mean() if tv_terms else l2 * 0.0 + return l2 + float(tv_weight) * tv + + +class DenseUVParserLoss(nn.Module): + def __init__( + self, + lambda_foreground=1.0, + lambda_layer=1.0, + lambda_part=0.5, + lambda_face=0.5, + lambda_layer_face=1.0, + lambda_uv=0.25, + lambda_uv_class=1.0, + lambda_affine=1.0, + lambda_surface=1.0, + lambda_outer_false_positive=0.75, + lambda_outer_false_negative=0.75, + lambda_route_confidence=0.25, + lambda_primary_route_swap=1.0, + lambda_route_texel_consistency=0.25, + lambda_route_prior_regularization=0.001, + outer_false_positive_gamma=2.0, + outer_false_negative_gamma=2.0, + primary_route_swap_gamma=2.0, + route_prior_tv_weight=1.0, + route_class_weight_floor=0.75, + route_outer_class_weight_cap=1.0, + uv_size=64, + foreground_pos_weight_max=20.0, + use_uv=True, + affine_translation_limit=0.06, + affine_log_scale_limit=0.03, + ): + super().__init__() + self.lambda_foreground = lambda_foreground + self.lambda_layer = lambda_layer + self.lambda_part = lambda_part + self.lambda_face = lambda_face + self.lambda_layer_face = lambda_layer_face + self.lambda_uv = lambda_uv + self.lambda_uv_class = lambda_uv_class + self.lambda_affine = lambda_affine + self.lambda_surface = lambda_surface + self.lambda_outer_false_positive = lambda_outer_false_positive + self.lambda_outer_false_negative = lambda_outer_false_negative + self.lambda_route_confidence = float(lambda_route_confidence) + self.lambda_primary_route_swap = float(lambda_primary_route_swap) + self.lambda_route_texel_consistency = float( + lambda_route_texel_consistency + ) + self.lambda_route_prior_regularization = float( + lambda_route_prior_regularization + ) + self.outer_false_positive_gamma = outer_false_positive_gamma + self.outer_false_negative_gamma = outer_false_negative_gamma + self.primary_route_swap_gamma = float(primary_route_swap_gamma) + self.route_prior_tv_weight = float(route_prior_tv_weight) + self.route_class_weight_floor = route_class_weight_floor + self.route_outer_class_weight_cap = route_outer_class_weight_cap + self.uv_size = uv_size + self.foreground_pos_weight_max = foreground_pos_weight_max + self.use_uv = bool(use_uv) + self.affine_translation_limit = max(float(affine_translation_limit), 1e-6) + self.affine_log_scale_limit = max(float(affine_log_scale_limit), 1e-6) + + def forward(self, outputs, targets): + fg_target = targets["foreground"].float() + pos_count = fg_target.sum().clamp_min(1.0) + neg_count = (fg_target.numel() - fg_target.sum()).clamp_min(1.0) + pos_weight = (neg_count / pos_count).clamp(max=self.foreground_pos_weight_max) + loss_foreground_bce = F.binary_cross_entropy_with_logits( + outputs["foreground"], + fg_target, + pos_weight=pos_weight, + ) + fg_prob = torch.sigmoid(outputs["foreground"]) + dice_den = fg_prob.sum() + fg_target.sum() + loss_foreground_dice = 1.0 - (2.0 * (fg_prob * fg_target).sum() + 1.0) / (dice_den + 1.0) + loss_foreground = loss_foreground_bce + loss_foreground_dice + + zero = outputs["foreground"].new_tensor(0.0) + geometry_route_roles = outputs["layer"].shape[1] == 3 and "route_role" in targets + layer_target = targets["route_role"] if geometry_route_roles else targets["layer"] + loss_layer = ( + _balanced_cross_entropy( + outputs["layer"], + layer_target, + min_weight=self.route_class_weight_floor, + class_weight_caps=( + float("inf"), + self.route_outer_class_weight_cap, + float("inf"), + ), + ) + if geometry_route_roles + else _deterministic_cross_entropy( + outputs["layer"], + layer_target, + ignore_index=IGNORE_INDEX, + ) + ) + loss_primary_route_swap = ( + primary_route_swap_loss( + outputs["layer"], + layer_target, + gamma=self.primary_route_swap_gamma, + ) + if geometry_route_roles + else zero + ) + loss_route_texel_consistency = ( + projected_texel_consistency_loss( + outputs["layer"], + layer_target, + targets["uv"], + uv_size=self.uv_size, + ) + if geometry_route_roles + else zero + ) + loss_route_prior_regularization = route_prior_regularization( + outputs, tv_weight=self.route_prior_tv_weight + ) + loss_outer_false_positive = ( + outer_false_positive_loss( + outputs["layer"], + layer_target, + gamma=self.outer_false_positive_gamma, + ) + if geometry_route_roles + else zero + ) + weighted_outer_false_positive = ( + self.lambda_outer_false_positive * loss_outer_false_positive + ) + loss_outer_false_negative = ( + outer_false_negative_loss( + outputs["layer"], + layer_target, + gamma=self.outer_false_negative_gamma, + ) + if geometry_route_roles + else zero + ) + weighted_outer_false_negative = ( + self.lambda_outer_false_negative * loss_outer_false_negative + ) + loss_part = ( + _deterministic_cross_entropy( + outputs["part"], + targets["part"], + ignore_index=IGNORE_INDEX, + ) + if "part" in outputs + else zero + ) + loss_face = ( + _deterministic_cross_entropy( + outputs["face"], + targets["face"], + ignore_index=IGNORE_INDEX, + ) + if "face" in outputs + else zero + ) + layer_face_target = combine_layer_face(targets["layer"], targets["face"]) + if "layer_face" in outputs: + loss_layer_face = _balanced_cross_entropy(outputs["layer_face"], layer_face_target) + else: + loss_layer_face = zero + + fg_mask = targets["layer"] != IGNORE_INDEX + use_uv = self.use_uv and "uv" in outputs + if use_uv and fg_mask.any(): + pred_uv_px = outputs["uv"].permute(0, 2, 3, 1)[fg_mask] * (self.uv_size - 1) + target_uv_px = targets["uv"].permute(0, 2, 3, 1)[fg_mask] * (self.uv_size - 1) + loss_uv = F.smooth_l1_loss(pred_uv_px, target_uv_px) + loss_uv_l1_px = (pred_uv_px - target_uv_px).abs().mean() + else: + loss_uv = zero + loss_uv_l1_px = zero + + if use_uv and "uv_x" in outputs and "uv_y" in outputs and fg_mask.any(): + target_x, target_y = uv_class_targets(targets["uv"], targets["layer"], self.uv_size) + loss_uv_x = _deterministic_cross_entropy( + outputs["uv_x"], + target_x, + ignore_index=IGNORE_INDEX, + ) + loss_uv_y = _deterministic_cross_entropy( + outputs["uv_y"], + target_y, + ignore_index=IGNORE_INDEX, + ) + loss_uv_class = 0.5 * (loss_uv_x + loss_uv_y) + else: + loss_uv_x = zero + loss_uv_y = zero + loss_uv_class = zero + + if "affine" in outputs and "affine" in targets: + affine_error = outputs["affine"] - targets["affine"].to(outputs["affine"].dtype) + loss_affine_translation = F.smooth_l1_loss( + affine_error[:, :2] / self.affine_translation_limit, + torch.zeros_like(affine_error[:, :2]), + ) + loss_affine_scale = F.smooth_l1_loss( + affine_error[:, 2] / self.affine_log_scale_limit, + torch.zeros_like(affine_error[:, 2]), + ) + loss_affine = 0.5 * (loss_affine_translation + loss_affine_scale) + H, W = outputs["foreground"].shape[-2:] + err_affine_translation_px = 0.5 * ( + affine_error[:, 0].abs().mean() * (W / 2.0) + + affine_error[:, 1].abs().mean() * (H / 2.0) + ) + err_affine_scale_pct = ( + outputs["affine"][:, 2].exp() - targets["affine"][:, 2].to(outputs["affine"].dtype).exp() + ).abs().mean() * 100.0 + else: + loss_affine_translation = zero + loss_affine_scale = zero + loss_affine = zero + err_affine_translation_px = zero + err_affine_scale_pct = zero + + if "surface" in outputs and "surface" in targets: + loss_surface = _balanced_cross_entropy(outputs["surface"], targets["surface"]) + acc_surface = _masked_accuracy(outputs["surface"], targets["surface"]) + else: + loss_surface = zero + acc_surface = zero + + if "route_confidence" in outputs and "route_role" in targets: + valid_route = targets["route_role"] != IGNORE_INDEX + predicted_role = outputs["layer"].detach().argmax(dim=1) + route_correct = predicted_role == targets["route_role"] + if "surface" in outputs and "surface" in targets: + predicted_surface = outputs["surface"].detach().argmax(dim=1) + secondary = targets["route_role"] == 2 + surface_correct = predicted_surface == targets["surface"] + route_correct = route_correct & (~secondary | surface_correct) + confidence_target = route_correct.float() + confidence_logits = outputs["route_confidence"][:, 0].float() + if valid_route.any(): + loss_route_confidence = F.binary_cross_entropy_with_logits( + confidence_logits[valid_route], + confidence_target[valid_route], + ) + confidence_probability = torch.sigmoid(confidence_logits) + confidence_mae = ( + confidence_probability[valid_route] + - confidence_target[valid_route] + ).abs().mean() + trusted_prediction = confidence_probability[valid_route] >= 0.5 + trusted_correct = confidence_target[valid_route] > 0.5 + precision_trusted_route = ( + (trusted_prediction & trusted_correct).float().sum() + / trusted_prediction.float().sum().clamp_min(1.0) + ) + coverage_trusted_route = trusted_prediction.float().mean() + else: + loss_route_confidence = zero + confidence_mae = zero + precision_trusted_route = zero + coverage_trusted_route = zero + else: + loss_route_confidence = zero + confidence_mae = zero + precision_trusted_route = zero + coverage_trusted_route = zero + + loss_routing = ( + loss_layer + + loss_affine + + loss_surface + + weighted_outer_false_positive + + weighted_outer_false_negative + + self.lambda_primary_route_swap * loss_primary_route_swap + + self.lambda_route_texel_consistency + * loss_route_texel_consistency + + self.lambda_route_prior_regularization + * loss_route_prior_regularization + if geometry_route_roles + else loss_surface + loss_uv_class + loss_layer_face + ) + loss_geometry = loss_foreground + loss_layer + loss_affine + ( + loss_surface if geometry_route_roles else zero + ) + weighted_outer_false_positive + weighted_outer_false_negative + ( + self.lambda_primary_route_swap * loss_primary_route_swap + + self.lambda_route_texel_consistency + * loss_route_texel_consistency + + self.lambda_route_prior_regularization + * loss_route_prior_regularization + if geometry_route_roles + else zero + ) + + loss_total = ( + self.lambda_foreground * loss_foreground + + self.lambda_layer * loss_layer + + self.lambda_part * loss_part + + self.lambda_face * loss_face + + self.lambda_layer_face * loss_layer_face + + self.lambda_uv * loss_uv + + self.lambda_uv_class * loss_uv_class + + self.lambda_affine * loss_affine + + self.lambda_surface * loss_surface + + weighted_outer_false_positive + + weighted_outer_false_negative + + self.lambda_route_confidence * loss_route_confidence + + self.lambda_primary_route_swap * loss_primary_route_swap + + self.lambda_route_texel_consistency + * loss_route_texel_consistency + + self.lambda_route_prior_regularization + * loss_route_prior_regularization + ) + + metrics = { + "loss_total": loss_total, + "loss_foreground": loss_foreground, + "loss_foreground_bce": loss_foreground_bce, + "loss_foreground_dice": loss_foreground_dice, + "loss_layer": loss_layer, + "loss_part": loss_part, + "loss_face": loss_face, + "loss_layer_face": loss_layer_face, + "loss_uv": loss_uv, + "loss_uv_l1_px": loss_uv_l1_px, + "loss_uv_class": loss_uv_class, + "loss_uv_x": loss_uv_x, + "loss_uv_y": loss_uv_y, + "loss_affine": loss_affine, + "loss_affine_translation": loss_affine_translation, + "loss_affine_scale": loss_affine_scale, + "err_affine_translation_px": err_affine_translation_px, + "err_affine_scale_pct": err_affine_scale_pct, + "loss_surface": loss_surface, + "loss_outer_false_positive": loss_outer_false_positive, + "loss_outer_false_positive_weighted": weighted_outer_false_positive, + "loss_outer_false_negative": loss_outer_false_negative, + "loss_outer_false_negative_weighted": weighted_outer_false_negative, + "loss_route_confidence": loss_route_confidence, + "loss_primary_route_swap": loss_primary_route_swap, + "loss_route_texel_consistency": loss_route_texel_consistency, + "loss_route_prior_regularization": loss_route_prior_regularization, + "confidence_mae": confidence_mae, + "precision_trusted_route": precision_trusted_route, + "coverage_trusted_route": coverage_trusted_route, + "loss_routing": loss_routing, + "loss_geometry": loss_geometry, + "acc_surface": acc_surface, + } + metrics.update(classification_metrics(outputs, targets, self.uv_size, use_uv=use_uv)) + return metrics + + +def uv_class_targets(uv, layer, uv_size): + valid = layer != IGNORE_INDEX + xy = (uv * (uv_size - 1)).round().long().clamp(0, uv_size - 1) + target_x = torch.full_like(layer, IGNORE_INDEX) + target_y = torch.full_like(layer, IGNORE_INDEX) + target_x[valid] = xy[:, 0][valid] + target_y[valid] = xy[:, 1][valid] + return target_x, target_y + + +def _masked_accuracy(logits, target): + mask = target != IGNORE_INDEX + if not mask.any(): + return logits.new_tensor(0.0) + pred = logits.argmax(dim=1) + return (pred[mask] == target[mask]).float().mean() + + +def classification_metrics(outputs, targets, uv_size, use_uv=True): + fg_pred = torch.sigmoid(outputs["foreground"]) > 0.5 + fg_target = targets["foreground"] > 0.5 + fg_acc = (fg_pred == fg_target).float().mean() + fg_pred_flat = fg_pred[:, 0] + fg_target_flat = fg_target[:, 0] + tp = (fg_pred_flat & fg_target_flat).sum().float() + fp = (fg_pred_flat & ~fg_target_flat).sum().float() + fn = (~fg_pred_flat & fg_target_flat).sum().float() + geometry_route_roles = outputs["layer"].shape[1] == 3 and "route_role" in targets + layer_target = targets["route_role"] if geometry_route_roles else targets["layer"] + metrics = { + "acc_foreground": fg_acc, + "precision_foreground": tp / (tp + fp).clamp_min(1.0), + "recall_foreground": tp / (tp + fn).clamp_min(1.0), + "iou_foreground": tp / (tp + fp + fn).clamp_min(1.0), + "acc_layer": _masked_accuracy(outputs["layer"], layer_target), + } + if geometry_route_roles: + metrics["acc_route_role"] = metrics["acc_layer"] + valid_role = layer_target != IGNORE_INDEX + pred_role = outputs["layer"].argmax(dim=1) + for name, role in (("inner", 0), ("outer", 1), ("secondary", 2)): + target_role = layer_target == role + predicted_role = pred_role == role + role_tp = (predicted_role & target_role & valid_role).sum().float() + role_fp = (predicted_role & ~target_role & valid_role).sum().float() + role_fn = (~predicted_role & target_role & valid_role).sum().float() + metrics[f"precision_{name}"] = role_tp / (role_tp + role_fp).clamp_min(1.0) + metrics[f"recall_{name}"] = role_tp / (role_tp + role_fn).clamp_min(1.0) + metrics[f"iou_{name}"] = role_tp / (role_tp + role_fp + role_fn).clamp_min(1.0) + if name in ("inner", "outer"): + metrics[f"count_{name}_tp"] = role_tp + metrics[f"count_{name}_fp"] = role_fp + metrics[f"count_{name}_fn"] = role_fn + valid_layer = torch.zeros_like(valid_role) + else: + valid_layer = targets["layer"] != IGNORE_INDEX + if valid_layer.any(): + pred_outer = outputs["layer"].argmax(dim=1) == 1 + target_outer = targets["layer"] == 1 + outer_tp = (pred_outer & target_outer & valid_layer).sum().float() + outer_fp = (pred_outer & ~target_outer & valid_layer).sum().float() + outer_fn = (~pred_outer & target_outer & valid_layer).sum().float() + metrics.update( + { + "precision_outer": outer_tp / (outer_tp + outer_fp).clamp_min(1.0), + "recall_outer": outer_tp / (outer_tp + outer_fn).clamp_min(1.0), + "iou_outer": outer_tp / (outer_tp + outer_fp + outer_fn).clamp_min(1.0), + "count_outer_tp": outer_tp, + "count_outer_fp": outer_fp, + "count_outer_fn": outer_fn, + } + ) + if "part" in outputs: + metrics["acc_part"] = _masked_accuracy(outputs["part"], targets["part"]) + if "face" in outputs: + metrics["acc_face"] = _masked_accuracy(outputs["face"], targets["face"]) + if "layer_face" in outputs: + metrics["acc_layer_face"] = _masked_accuracy( + outputs["layer_face"], + combine_layer_face(targets["layer"], targets["face"]), + ) + if use_uv and "uv_x" in outputs and "uv_y" in outputs: + target_x, target_y = uv_class_targets(targets["uv"], targets["layer"], uv_size) + valid = target_x != IGNORE_INDEX + if valid.any(): + pred_x = outputs["uv_x"].argmax(dim=1) + pred_y = outputs["uv_y"].argmax(dim=1) + err_x = (pred_x[valid] - target_x[valid]).abs() + err_y = (pred_y[valid] - target_y[valid]).abs() + metrics.update( + { + "acc_uv_x": (pred_x[valid] == target_x[valid]).float().mean(), + "acc_uv_y": (pred_y[valid] == target_y[valid]).float().mean(), + "acc_uv_exact": ((pred_x[valid] == target_x[valid]) & (pred_y[valid] == target_y[valid])).float().mean(), + "acc_uv_within1": ((err_x <= 1) & (err_y <= 1)).float().mean(), + "err_uv_class_l1_px": 0.5 * (err_x.float().mean() + err_y.float().mean()), + } + ) + return metrics diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/material_refine.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/material_refine.py new file mode 100644 index 0000000000000000000000000000000000000000..649d80dc97a64c8415bb89ca36d066a1c4c22be2 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/material_refine.py @@ -0,0 +1,108 @@ +"""Fit visible head materials while keeping ownership and hidden layers isolated.""" +import torch +from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology +from SkingToolkit.dense_uv_parser.accessories import head_bounds + + +def refine_head_material(uv, images, sources, renderer, views, steps=48, ownership=None, inner_exclusion=None, inner_texel_support=None, headwear_layer=None, headwear_uv=None, headwear_family=None): + original=uv.detach().float() + topology=build_simple_uv_topology() + head=(topology.valid&(topology.part==0)).to(uv.device)[None,None] + inner=head&(topology.layer.to(uv.device)[None,None]==0) + outer=head&~inner + trainable_inner=inner if inner_texel_support is None else inner&inner_texel_support[:,None] + valid=torch.zeros_like(sources,dtype=torch.bool) + y0,y1,x0,x1=head_bounds(*sources.shape[-2:]) + valid[:,y0:y1,x0:x1]=sources[:,y0:y1,x0:x1] + if not valid.any():return original + target=images[:,:3].detach().float() + def render(skin): + return torch.stack([renderer.forward_view(skin,v) for v in views],1).flatten(0,1)[:,:3] + # At fixed alpha the renderer is linear in RGB. Unit-colour probes measure + # actual contributions, including its secondary surfaces and bilinear edges. + with torch.no_grad(): + zero=torch.cat([torch.zeros_like(original[:,:3]),original[:,3:4]],1) + base=render(zero) + weights=[] + for surface in (inner,outer): + probe=torch.cat([surface.expand_as(original[:,:3]).float(),original[:,3:4]],1) + contribution=(render(probe)-base).mean(1).clamp(0,1) + weights.append((valid&(contribution>=.98))[:,None].float()) + if ownership is not None: + probability,identity=ownership.float().softmax(1).max(1) + accessory=((identity==3)|(identity==4))&(probability>=.5) + weights[0]=weights[0]*(~accessory[:,None]) + if inner_exclusion is not None: + weights[0]=weights[0]*(~inner_exclusion[:,None]) + if headwear_layer is not None: + for layer in (0,1): + weights[layer] *= ((headwear_layer < 0) | (headwear_layer == layer))[:,None] + if headwear_uv is not None and headwear_family is not None: + # Only compare like materials. The renderer must see the SAME + # semantic component as the source pixel before its RGB may fit it. + compatible = torch.zeros_like(valid) + for family in range(5): + surface = (headwear_uv[:,None] == family) & head + probe = torch.cat([surface.expand_as(original[:,:3]).float(), original[:,3:4]],1) + contribution = (render(probe)-base).mean(1) + compatible |= (headwear_family == family) & (contribution >= .98) + for layer in (0,1): + weights[layer] *= compatible[:,None] + denominators=[(w.sum()*3).clamp_min(1) for w in weights] + if not any(w.any() for w in weights):return original + def losses(skin): + error=(render(skin)-target).square() + return [(error*w).sum()/den for w,den in zip(weights,denominators)] + with torch.enable_grad(): + rgb=original[:,:3].clone().requires_grad_(True) + optimizer=torch.optim.Adam([rgb],lr=.06) + with torch.no_grad():best_error=float(sum(losses(original)));best=original.clone() + for _ in range(steps): + optimizer.zero_grad(set_to_none=True) + skin=torch.cat([torch.where(head,rgb,original[:,:3]),original[:,3:4]],1) + loss_inner,loss_outer=losses(skin) + if not torch.isfinite(loss_inner+loss_outer):raise RuntimeError('Non-finite head material objective') + gi=torch.autograd.grad(loss_inner,rgb,retain_graph=True)[0] + go=torch.autograd.grad(loss_outer,rgb)[0] + rgb.grad=torch.where(trainable_inner,gi,torch.where(outer,go,0.)) + optimizer.step() + with torch.no_grad(): + rgb.clamp_(0,1) + candidate=torch.cat([torch.where(head,rgb,original[:,:3]),original[:,3:4]],1) + error=float(sum(losses(candidate))) + if error.9 + anchor=rgba[:3,opaque].mean(1) if opaque.any() else uniform((3,)) + bg=(anchor[:,None,None]+uniform((3,1,1),-.12,.12)).clamp(0,1).expand(3,size,size).clone() + else: + bg=(F.interpolate(uniform((1,3,4,4)),(size,size),mode='bilinear',align_corners=False)[0]+uniform((3,size,size),-.06,.06)).clamp(0,1) + if stress: + # Modify foreground illumination before compositing, keeping true alpha. + tint=uniform((3,1,1),.85,1.15) + premult=(premult*tint).clamp_min(0).minimum(alpha) + image=compose(premult,alpha,bg) + if stress and rng.random()<.35: + buffer=io.BytesIO();Image.fromarray((image.permute(1,2,0).numpy()*255).round().astype('uint8')).save(buffer,format='JPEG',quality=int(rng.integers(65,96))) + image=torch.from_numpy(np.array(Image.open(io.BytesIO(buffer.getvalue())).convert('RGB')).copy()).permute(2,0,1).float()/255 + elif stress and rng.random()<.25: + image=(image+uniform(tuple(image.shape),-.009,.009)).clamp(0,1) + return image,alpha,kind + + +class MattingDataset(Dataset): + def __init__(self,manifest,split,size=512,limit=None,epoch=0): + content=json.loads(Path(manifest).read_text()) + if 'canonical_identity_check' not in content:raise ValueError('Dataset must pass normalized skin identity deduplication') + self.records=[r for r in content['records'] if r['split']==split] + if limit:self.records=self.records[:limit] + self.size,self.epoch,self.split=size,epoch,split + self.offset={'train':710100,'validation':1710100,'test':2710100}[split] + def __len__(self):return len(self.records) + def __getitem__(self,index): + rgba=torch.from_numpy(np.array(Image.open(self.records[index]['file']).convert('RGBA')).copy()).permute(2,0,1).float()/255 + image,alpha,kind=make_composite(rgba,self.size,self.offset+index*104729+self.epoch*1000003) + return {'image':image,'alpha':alpha,'background_kind':kind} diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/model.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/model.py new file mode 100644 index 0000000000000000000000000000000000000000..a12f1629337c3cf7758335179a2ff8cf0454f2ad --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/model.py @@ -0,0 +1,816 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from SkingToolkit.dense_uv_parser.uv_topology import ( + build_outer_uv_graph, + build_simple_uv_topology, +) + + +def norm_groups(channels): + for groups in (8, 4, 2, 1): + if channels % groups == 0: + return groups + return 1 + + +class ConvBlock(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + nn.GroupNorm(norm_groups(out_channels), out_channels), + nn.SiLU(inplace=True), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + nn.GroupNorm(norm_groups(out_channels), out_channels), + nn.SiLU(inplace=True), + ) + + def forward(self, x): + return self.block(x) + + +class DownBlock(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.down = nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1) + self.block = ConvBlock(out_channels, out_channels) + + def forward(self, x): + return self.block(self.down(x)) + + +class UpBlock(nn.Module): + def __init__(self, in_channels, skip_channels, out_channels): + super().__init__() + self.block = ConvBlock(in_channels + skip_channels, out_channels) + + def forward(self, x, skip): + x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) + return self.block(torch.cat([x, skip], dim=1)) + + +class MultiViewSemanticFusion(nn.Module): + """Fuse frozen per-view semantic features into parser bottleneck features.""" + + def __init__( + self, + raw_feature_dim, + semantic_channels, + bottleneck_channels, + view_classes, + attention_heads=4, + layers=1, + dropout=0.05, + ): + super().__init__() + if raw_feature_dim < 1 or semantic_channels < 1: + raise ValueError("Semantic feature dimensions must be positive.") + if view_classes < 1: + raise ValueError("Semantic fusion requires view-conditioned parser inputs.") + if attention_heads < 1: + raise ValueError("semantic_attention_heads must be positive.") + if semantic_channels % attention_heads != 0: + raise ValueError("semantic_channels must be divisible by attention_heads.") + if layers < 1: + raise ValueError("semantic_layers must be positive.") + if not 0.0 <= dropout < 1.0: + raise ValueError("semantic_dropout must be in [0, 1).") + self.raw_feature_dim = int(raw_feature_dim) + self.semantic_channels = int(semantic_channels) + self.view_classes = int(view_classes) + self.input_projection = nn.Sequential( + nn.LayerNorm(raw_feature_dim), + nn.Linear(raw_feature_dim, semantic_channels), + nn.GELU(), + ) + self.view_embedding = nn.Parameter( + torch.randn(view_classes, semantic_channels) * 0.02 + ) + self.encoder = nn.ModuleList( + nn.TransformerEncoderLayer( + d_model=semantic_channels, + nhead=attention_heads, + dim_feedforward=semantic_channels * 4, + dropout=dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.modulation = nn.Sequential( + nn.LayerNorm(semantic_channels * 2), + nn.Linear(semantic_channels * 2, semantic_channels * 2), + nn.GELU(), + nn.Linear(semantic_channels * 2, bottleneck_channels * 2), + ) + self.summary = nn.Sequential( + nn.LayerNorm(semantic_channels), + nn.Linear(semantic_channels, semantic_channels), + nn.GELU(), + ) + # Preserve the geometry-only initialization. The adapter starts by + # contributing no FiLM shift/scale and learns semantic corrections. + nn.init.zeros_(self.modulation[-1].weight) + nn.init.zeros_(self.modulation[-1].bias) + + def forward(self, raw_features, view_ids, sample_count): + if raw_features.dim() == 3: + # Training may cache multiple complete view groups per skin, for + # example the inference pair followed by a privileged right-side + # pair. Each group is still fused independently below. + raw_features = raw_features.reshape(-1, raw_features.shape[-1]) + if raw_features.dim() != 2 or raw_features.shape != ( + sample_count, + self.raw_feature_dim, + ): + raise ValueError( + "Semantic features must be shaped NxD or BxVxD; got " + f"{tuple(raw_features.shape)}." + ) + if sample_count % self.view_classes != 0: + raise ValueError( + f"Semantic sample count {sample_count} is not divisible by " + f"view count {self.view_classes}." + ) + grouped_view_ids = view_ids.reshape(-1, self.view_classes) + expected_ids = torch.arange( + self.view_classes, device=view_ids.device + ).view(1, -1) + if not torch.equal(grouped_view_ids, expected_ids.expand_as(grouped_view_ids)): + raise ValueError("Semantic fusion requires canonical grouped view order.") + + batch = sample_count // self.view_classes + tokens = self.input_projection(raw_features.float()).reshape( + batch, self.view_classes, self.semantic_channels + ) + tokens = tokens + self.view_embedding.unsqueeze(0) + for layer in self.encoder: + tokens = layer(tokens) + pooled = tokens.mean(dim=1) + per_view = torch.cat( + [tokens, pooled.unsqueeze(1).expand_as(tokens)], dim=-1 + ).reshape(sample_count, self.semantic_channels * 2) + return self.modulation(per_view), self.summary(pooled) + + +class SpatialSemanticFusion(nn.Module): + """Project frozen 2D semantic features into the parser bottleneck. + + The final projection starts at zero, preserving the stable geometry parser + at initialization while allowing spatial semantic corrections to emerge + during training. + """ + + def __init__( + self, + raw_feature_dim, + semantic_channels, + bottleneck_channels, + ): + super().__init__() + if raw_feature_dim < 1 or semantic_channels < 1: + raise ValueError("Spatial semantic feature dimensions must be positive.") + self.raw_feature_dim = int(raw_feature_dim) + self.semantic_channels = int(semantic_channels) + self.input_norm = nn.LayerNorm(raw_feature_dim) + self.input_projection = nn.Conv2d( + raw_feature_dim, semantic_channels, kernel_size=1 + ) + self.activation = nn.GELU() + self.output_projection = nn.Conv2d( + semantic_channels, bottleneck_channels, kernel_size=1 + ) + nn.init.zeros_(self.output_projection.weight) + nn.init.zeros_(self.output_projection.bias) + + def forward(self, raw_features, sample_count, output_size): + if raw_features.dim() == 5: + raw_features = raw_features.reshape( + -1, + raw_features.shape[-3], + raw_features.shape[-2], + raw_features.shape[-1], + ) + if ( + raw_features.dim() != 4 + or raw_features.shape[0] != sample_count + or raw_features.shape[1] != self.raw_feature_dim + ): + raise ValueError( + "Spatial semantic features must be shaped NxCxHxW or " + f"BxVxCxHxW; got {tuple(raw_features.shape)}." + ) + normalized = self.input_norm( + raw_features.float().permute(0, 2, 3, 1) + ).permute(0, 3, 1, 2) + residual = self.output_projection( + self.activation(self.input_projection(normalized)) + ) + return F.interpolate( + residual, + size=output_size, + mode="bilinear", + align_corners=False, + ) + + +class OuterUVGraphBlock(nn.Module): + """Message passing over physical neighbours of Minecraft outer texels.""" + + def __init__(self, channels, dropout=0.05): + super().__init__() + self.self_projection = nn.Linear(channels, channels) + self.neighbour_projection = nn.Linear(channels, channels) + self.norm = nn.LayerNorm(channels) + self.dropout = nn.Dropout(dropout) + + def forward(self, nodes, edge_index, degree): + source, target = edge_index + neighbour_sum = torch.zeros_like(nodes) + neighbour_sum.index_add_(1, target, nodes[:, source]) + neighbour_mean = neighbour_sum / degree.view(1, -1, 1).clamp_min(1.0) + update = self.self_projection(nodes) + update = update + self.neighbour_projection(neighbour_mean) + return self.norm(nodes + self.dropout(F.gelu(update))) + + +class ProjectedOuterUVTopologyHead(nn.Module): + """Predict outer alpha from projected image features and cube topology.""" + + def __init__( + self, + input_channels, + global_context_dim, + hidden_channels=64, + layers=3, + dropout=0.05, + uv_size=64, + ): + super().__init__() + if layers < 1: + raise ValueError("outer_uv_topology_layers must be positive.") + flat_indices, edge_index = build_outer_uv_graph() + topology = build_simple_uv_topology() + parts = topology.part.reshape(-1)[flat_indices] + faces = topology.face.reshape(-1)[flat_indices] + local_uv = topology.local_uv.reshape(-1, 2)[flat_indices] + degree = torch.bincount( + edge_index[1], + minlength=flat_indices.numel(), + ).float() + self.uv_size = int(uv_size) + self.register_buffer("flat_indices", flat_indices, persistent=False) + self.register_buffer("edge_index", edge_index, persistent=False) + self.register_buffer("degree", degree, persistent=False) + self.register_buffer("parts", parts, persistent=False) + self.register_buffer("faces", faces, persistent=False) + self.register_buffer("local_uv", local_uv, persistent=False) + self.input_projection = nn.Linear(input_channels, hidden_channels) + self.part_embedding = nn.Embedding(6, hidden_channels) + self.face_embedding = nn.Embedding(6, hidden_channels) + self.position_projection = nn.Linear(2, hidden_channels) + self.global_projection = nn.Sequential( + nn.LayerNorm(global_context_dim), + nn.Linear(global_context_dim, hidden_channels), + nn.GELU(), + ) + self.blocks = nn.ModuleList( + OuterUVGraphBlock(hidden_channels, dropout=dropout) + for _ in range(layers) + ) + self.output = nn.Sequential( + nn.LayerNorm(hidden_channels), + nn.Linear(hidden_channels, 1), + ) + # A very negative initialization made the one-epoch recipe prone to + # the all-transparent solution after precision penalties were added. + # Keep a sparse prior without starving early positive gradients. + nn.init.constant_(self.output[-1].bias, -1.0) + + def forward(self, atlas_features, global_context): + if atlas_features.dim() != 4: + raise ValueError("Projected outer UV features must be BCHW.") + if atlas_features.shape[-2:] != (self.uv_size, self.uv_size): + raise ValueError( + "Projected outer UV feature size must match the UV atlas." + ) + flattened = atlas_features.flatten(2).transpose(1, 2) + nodes = flattened.index_select(1, self.flat_indices) + nodes = self.input_projection(nodes.float()) + nodes = nodes + self.part_embedding(self.parts).unsqueeze(0) + nodes = nodes + self.face_embedding(self.faces).unsqueeze(0) + nodes = nodes + self.position_projection(self.local_uv).unsqueeze(0) + nodes = nodes + self.global_projection(global_context.float()).unsqueeze(1) + for block in self.blocks: + nodes = block(nodes, self.edge_index, self.degree) + node_logits = self.output(nodes).squeeze(-1) + atlas_logits = node_logits.new_full( + (node_logits.shape[0], self.uv_size * self.uv_size), + -12.0, + ) + atlas_logits[:, self.flat_indices] = node_logits + return atlas_logits.reshape( + node_logits.shape[0], + 1, + self.uv_size, + self.uv_size, + ) + + +class DenseUVParserNet(nn.Module): + """Predict dense Minecraft UV routing for each render pixel.""" + + def __init__( + self, + input_channels=4, + base_channels=32, + part_classes=6, + face_classes=6, + layer_classes=None, + layer_face_classes=12, + uv_size=64, + uv_classification=True, + view_classes=0, + predict_affine=False, + affine_translation_scale=0.0, + affine_scale_range=0.0, + surface_classes=0, + geometry_only=False, + feature_dropout=0.0, + semantic_feature_dim=0, + semantic_channels=128, + semantic_attention_heads=4, + semantic_layers=1, + semantic_dropout=0.05, + semantic_spatial_feature_dim=0, + semantic_spatial_channels=64, + predict_confidence=False, + route_role_spatial_prior=False, + route_prior_height=32, + route_prior_width=16, + route_prior_logit_cap=1.5, + route_prior_dropout=0.10, + predict_head_accessories=False, + predict_hat_components=False, + predict_head_ownership=False, + predict_headwear=False, + predict_headwear_presence=False, + predict_headphone_presence=False, + accessory_route_threshold=0.90, + predict_outer_uv_occupancy=False, + outer_uv_feature_channels=32, + outer_uv_topology_channels=64, + outer_uv_topology_layers=3, + outer_uv_topology_dropout=0.05, + outer_uv_route_evidence_dropout=1.0, + ): + super().__init__() + self.accessory_route_threshold = float(accessory_route_threshold) + if not 0.5 <= self.accessory_route_threshold <= 1.: + raise ValueError("Accessory seed confidence must lie in [0.5, 1]") + self.predict_hat_components = bool(predict_hat_components) + self.predict_head_accessories = bool(predict_head_accessories) + if self.predict_head_accessories: + from SkingToolkit.dense_uv_parser.accessories import HeadAccessoryHead + self.accessory_head = HeadAccessoryHead(semantic_spatial_feature_dim, self.predict_hat_components) + self.predict_headphone_presence = bool(predict_headphone_presence) + self.predict_headwear_presence = bool(predict_headwear_presence) + self.predict_headwear = bool(predict_headwear) + if self.predict_headwear: + from SkingToolkit.dense_uv_parser.headwear import HeadwearHead + self.headwear_head = HeadwearHead(semantic_spatial_feature_dim,self.predict_headwear_presence) + self.predict_head_ownership = bool(predict_head_ownership) + if self.predict_head_ownership: + from SkingToolkit.dense_uv_parser.ownership import HeadOwnershipHead + self.ownership_head = HeadOwnershipHead(semantic_spatial_feature_dim,self.predict_headphone_presence) + self.geometry_only = bool(geometry_only) + if layer_classes is None: + layer_classes = 3 if self.geometry_only else 2 + self.layer_classes = int(layer_classes) + self.uv_size = int(uv_size) + self.uv_classification = bool(uv_classification) and not self.geometry_only + self.view_classes = int(view_classes) + self.predict_affine = bool(predict_affine) + self.affine_translation_scale = float(affine_translation_scale) + self.affine_scale_range = float(affine_scale_range) + self.surface_classes = int(surface_classes) + self.semantic_feature_dim = int(semantic_feature_dim) + self.semantic_channels = int(semantic_channels) + self.semantic_attention_heads = int(semantic_attention_heads) + self.semantic_layers = int(semantic_layers) + self.semantic_dropout = float(semantic_dropout) + self.semantic_spatial_feature_dim = int(semantic_spatial_feature_dim) + self.semantic_spatial_channels = int(semantic_spatial_channels) + self.predict_confidence = bool(predict_confidence) + self.route_role_spatial_prior = bool(route_role_spatial_prior) + self.route_prior_height = int(route_prior_height) + self.route_prior_width = int(route_prior_width) + self.route_prior_logit_cap = float(route_prior_logit_cap) + self.route_prior_dropout = float(route_prior_dropout) + self.predict_outer_uv_occupancy = bool(predict_outer_uv_occupancy) + self.outer_uv_feature_channels = int(outer_uv_feature_channels) + self.outer_uv_topology_channels = int(outer_uv_topology_channels) + self.outer_uv_topology_layers = int(outer_uv_topology_layers) + self.outer_uv_topology_dropout = float(outer_uv_topology_dropout) + self.outer_uv_route_evidence_dropout = float( + outer_uv_route_evidence_dropout + ) + if self.route_prior_height < 1 or self.route_prior_width < 1: + raise ValueError("Route-prior dimensions must be positive.") + if self.route_prior_logit_cap <= 0.0: + raise ValueError("route_prior_logit_cap must be positive.") + if not 0.0 <= self.route_prior_dropout < 1.0: + raise ValueError("route_prior_dropout must be in [0, 1).") + if self.route_role_spatial_prior and ( + not self.geometry_only or self.view_classes < 1 + ): + raise ValueError( + "The fixed-view route-role prior requires geometry_only with view classes." + ) + self.feature_dropout_probability = float(feature_dropout) + if not 0.0 <= self.feature_dropout_probability < 1.0: + raise ValueError("feature_dropout must be in [0, 1).") + if not 0.0 <= self.outer_uv_route_evidence_dropout <= 1.0: + raise ValueError( + "outer_uv_route_evidence_dropout must be in [0, 1]." + ) + c = base_channels + self.stem = ConvBlock(input_channels + self.view_classes, c) + self.down1 = DownBlock(c, c * 2) + self.down2 = DownBlock(c * 2, c * 4) + self.down3 = DownBlock(c * 4, c * 8) + self.mid = ConvBlock(c * 8, c * 8) + self.semantic_fusion = ( + MultiViewSemanticFusion( + self.semantic_feature_dim, + self.semantic_channels, + c * 8, + self.view_classes, + attention_heads=self.semantic_attention_heads, + layers=self.semantic_layers, + dropout=self.semantic_dropout, + ) + if self.semantic_feature_dim > 0 + else None + ) + self.semantic_spatial_fusion = ( + SpatialSemanticFusion( + self.semantic_spatial_feature_dim, + self.semantic_spatial_channels, + c * 8, + ) + if self.semantic_spatial_feature_dim > 0 + else None + ) + self.up2 = UpBlock(c * 8, c * 4, c * 4) + self.up1 = UpBlock(c * 4, c * 2, c * 2) + self.up0 = UpBlock(c * 2, c, c) + self.features = nn.Sequential( + nn.Conv2d(c, c, kernel_size=3, padding=1), + nn.SiLU(inplace=True), + ) + self.feature_dropout = nn.Dropout2d(self.feature_dropout_probability) + self.foreground = nn.Conv2d(c, 1, kernel_size=1) + self.layer = nn.Conv2d(c, self.layer_classes, kernel_size=1) + self.route_role_prior = ( + nn.Parameter( + torch.zeros( + self.view_classes, + self.layer_classes, + self.route_prior_height, + self.route_prior_width, + ) + ) + if self.route_role_spatial_prior + else None + ) + self.route_confidence = ( + nn.Conv2d(c, 1, kernel_size=1) if self.predict_confidence else None + ) + if self.semantic_fusion is not None: + self.outer_presence_head = nn.Linear(self.semantic_channels, 6) + self.outer_coverage_head = nn.Linear(self.semantic_channels, 6) + if self.predict_outer_uv_occupancy: + if self.view_classes < 1: + raise ValueError( + "Outer UV occupancy prediction requires grouped fixed-view inputs." + ) + if self.outer_uv_feature_channels < 1: + raise ValueError("outer_uv_feature_channels must be positive.") + occupancy_context_channels = c * 8 + ( + self.semantic_channels if self.semantic_fusion is not None else 0 + ) + self.outer_uv_feature_projection = nn.Sequential( + nn.Conv2d(c, self.outer_uv_feature_channels, kernel_size=1), + nn.GELU(), + ) + # Projected feature mean + mean/max p_outer + foreground coverage + # + number of supporting views. + occupancy_input_channels = self.outer_uv_feature_channels + 4 + self.outer_uv_occupancy_head = ProjectedOuterUVTopologyHead( + occupancy_input_channels, + occupancy_context_channels, + hidden_channels=self.outer_uv_topology_channels, + layers=self.outer_uv_topology_layers, + dropout=self.outer_uv_topology_dropout, + uv_size=uv_size, + ) + if not self.geometry_only: + self.part = nn.Conv2d(c, part_classes, kernel_size=1) + self.face = nn.Conv2d(c, face_classes, kernel_size=1) + self.layer_face = ( + nn.Conv2d(c, layer_face_classes, kernel_size=1) + if layer_face_classes > 0 + else None + ) + self.uv = nn.Conv2d(c, 2, kernel_size=1) + if self.uv_classification: + self.uv_x = nn.Conv2d(c, uv_size, kernel_size=1) + self.uv_y = nn.Conv2d(c, uv_size, kernel_size=1) + else: + self.layer_face = None + if self.predict_affine: + if not self.geometry_only and self.surface_classes < 2: + raise ValueError("Global-affine routing requires at least two static surface classes.") + if self.surface_classes == 1: + raise ValueError("Surface routing requires at least two static surface classes.") + if self.surface_classes > 0: + self.surface = nn.Conv2d(c, self.surface_classes, kernel_size=1) + hidden = max(c * 4, 32) + self.affine_head = nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.Linear(c * 8, hidden), + nn.SiLU(inplace=True), + nn.Linear(hidden, 3), + ) + # Start from the canonical pose; the residual head learns only the + # small global transform introduced by the configured augmentation. + nn.init.zeros_(self.affine_head[-1].weight) + nn.init.zeros_(self.affine_head[-1].bias) + + self.affine_translation_limit = 2.0 * self.affine_translation_scale + lower_log_scale = math.log(max(1.0 - self.affine_scale_range, 1e-6)) + upper_log_scale = math.log1p(self.affine_scale_range) + self.affine_log_scale_limit = max(abs(lower_log_scale), abs(upper_log_scale)) + + def _runtime_semantic_features(self, images, foreground=None): + backbone = getattr(self, "_runtime_semantic_backbone", None) + if backbone is None: + raise ValueError( + "This parser requires semantic_features or an attached semantic " + "runtime backbone." + ) + rgb = images[:, :3] + if foreground is not None: + if foreground.dim() == 3: + foreground = foreground.unsqueeze(1) + if ( + foreground.dim() != 4 + or foreground.shape[0] != rgb.shape[0] + or foreground.shape[1] != 1 + ): + raise ValueError( + "semantic_foreground must be shaped Nx1xHxW or NxHxW; got " + f"{tuple(foreground.shape)}." + ) + if foreground.shape[-2:] != rgb.shape[-2:]: + foreground = F.interpolate( + foreground.float(), + size=rgb.shape[-2:], + mode="nearest", + ) + foreground = foreground.to(device=rgb.device, dtype=rgb.dtype).clamp(0, 1) + # A fixed neutral background prevents the frozen spatial tower from + # learning accidental correlations with randomized training colors. + rgb = rgb * foreground + 0.5 * (1.0 - foreground) + with torch.no_grad(): + if hasattr(backbone, "encode_dense"): + return backbone.encode_dense(rgb) + return backbone.encode_global(rgb) + + def forward( + self, + x, + view_ids=None, + semantic_features=None, + semantic_foreground=None, + ): + source_images = x + if self.view_classes > 0: + if view_ids is None: + raise ValueError("view_ids are required for a view-conditioned dense UV parser.") + if view_ids.shape != (x.shape[0],): + raise ValueError(f"Expected view_ids shape {(x.shape[0],)}, got {tuple(view_ids.shape)}.") + if view_ids.min() < 0 or view_ids.max() >= self.view_classes: + raise ValueError(f"view_ids must be in [0, {self.view_classes - 1}].") + view_one_hot = F.one_hot(view_ids.long(), num_classes=self.view_classes).to(dtype=x.dtype) + view_one_hot = view_one_hot.view(x.shape[0], self.view_classes, 1, 1) + x = torch.cat([x, view_one_hot.expand(-1, -1, x.shape[2], x.shape[3])], dim=1) + + s0 = self.stem(x) + s1 = self.down1(s0) + s2 = self.down2(s1) + s3 = self.down3(s2) + x = self.mid(s3) + if ( + self.semantic_fusion is not None + or self.semantic_spatial_fusion is not None + ) and semantic_features is None: + semantic_features = self._runtime_semantic_features( + source_images, foreground=semantic_foreground + ) + if isinstance(semantic_features, dict): + semantic_global = semantic_features.get( + "raw_global", semantic_features.get("global") + ) + semantic_spatial = semantic_features.get( + "raw_spatial", semantic_features.get("spatial") + ) + else: + semantic_global = semantic_features + semantic_spatial = None + if self.semantic_spatial_fusion is not None: + if semantic_spatial is None: + raise ValueError( + "This parser checkpoint requires spatial semantic features." + ) + x = x + self.semantic_spatial_fusion( + semantic_spatial, + source_images.shape[0], + x.shape[-2:], + ).to(dtype=x.dtype) + visual_summary = x.mean(dim=(2, 3)) + if self.predict_outer_uv_occupancy: + if visual_summary.shape[0] % self.view_classes != 0: + raise ValueError( + "Outer UV occupancy prediction requires complete fixed-view groups." + ) + visual_summary = visual_summary.reshape( + -1, self.view_classes, visual_summary.shape[-1] + ).mean(dim=1) + semantic_summary = None + if self.semantic_fusion is not None: + if semantic_global is None: + raise ValueError( + "This parser checkpoint requires global semantic features." + ) + modulation, semantic_summary = self.semantic_fusion( + semantic_global, + view_ids, + source_images.shape[0], + ) + scale, shift = modulation.chunk(2, dim=1) + x = x * (1.0 + scale.unsqueeze(-1).unsqueeze(-1)) + x = x + shift.unsqueeze(-1).unsqueeze(-1) + affine = None + if self.predict_affine: + raw_affine = torch.tanh(self.affine_head(x)) + affine = torch.stack( + [ + raw_affine[:, 0] * self.affine_translation_limit, + raw_affine[:, 1] * self.affine_translation_limit, + raw_affine[:, 2] * self.affine_log_scale_limit, + ], + dim=1, + ) + x = self.up2(x, s2) + x = self.up1(x, s1) + x = self.up0(x, s0) + x = self.features(x) + occupancy_feature_source = x + x = self.feature_dropout(x) + layer_evidence = self.layer(x) + outputs = { + "foreground": self.foreground(x), + "layer": layer_evidence, + } + if self.route_role_prior is not None: + selected_prior_raw = self.route_role_prior.index_select( + 0, view_ids.long() + ) + selected_prior = self.route_prior_logit_cap * torch.tanh( + selected_prior_raw / self.route_prior_logit_cap + ) + if self.training and self.route_prior_dropout > 0.0: + keep = ( + torch.rand( + selected_prior.shape[0], + 1, + 1, + 1, + device=selected_prior.device, + ) + >= self.route_prior_dropout + ) + selected_prior = selected_prior * keep.to( + dtype=selected_prior.dtype + ) + route_prior = F.interpolate( + selected_prior, + size=layer_evidence.shape[-2:], + mode="bilinear", + align_corners=False, + ).to(dtype=layer_evidence.dtype) + outputs["layer"] = layer_evidence + route_prior + outputs["route_role_evidence"] = layer_evidence + outputs["route_role_prior"] = route_prior + outputs["route_role_prior_raw"] = self.route_role_prior + if self.route_confidence is not None: + outputs["route_confidence"] = self.route_confidence(x) + if semantic_summary is not None: + outputs["outer_presence_logits"] = self.outer_presence_head( + semantic_summary + ) + outputs["outer_coverage"] = torch.sigmoid( + self.outer_coverage_head(semantic_summary) + ) + if self.predict_outer_uv_occupancy: + occupancy_summary = ( + torch.cat([visual_summary, semantic_summary], dim=1) + if semantic_summary is not None + else visual_summary + ) + # Occupancy learns from the parser representation without allowing + # its sparse atlas loss to destabilize the proven image-space route + # trunk. A separate agreement loss later transfers calibrated UV + # structure back into the route head after warm-up. + outputs["outer_uv_features"] = self.outer_uv_feature_projection( + occupancy_feature_source.detach() + ) + outputs["outer_uv_global_context"] = occupancy_summary.detach() + if not self.geometry_only: + outputs["part"] = self.part(x) + outputs["face"] = self.face(x) + outputs["uv"] = torch.sigmoid(self.uv(x)) + if self.layer_face is not None: + outputs["layer_face"] = self.layer_face(x) + if self.uv_classification: + outputs["uv_x"] = self.uv_x(x) + outputs["uv_y"] = self.uv_y(x) + if affine is not None: + # [tx, ty, log_scale]. tx/ty are affine_grid normalized coordinates. + outputs["affine"] = affine + if self.surface_classes > 0: + outputs["surface"] = self.surface(x) + if self.predict_head_accessories: + prediction=self.predict_accessories(source_images,semantic_foreground,return_components=self.predict_hat_components) + if self.predict_hat_components: + outputs["accessory_logits"],outputs["hat_component_logits"]=prediction + else:outputs["accessory_logits"]=prediction + outputs["accessory_route_threshold"] = self.accessory_route_threshold + if self.predict_headwear: + prediction = self.predict_headwear_components(source_images, semantic_foreground,return_presence=self.predict_headwear_presence) + if self.predict_headwear_presence:outputs["headwear_logits"],outputs["headwear_presence_logits"]=prediction + else:outputs["headwear_logits"]=prediction + if self.predict_head_ownership: + prediction = self.predict_ownership(source_images, semantic_foreground, return_presence=self.predict_headphone_presence) + if self.predict_headphone_presence: + outputs["head_ownership_logits"],outputs["headphone_presence_logit"] = prediction + else:outputs["head_ownership_logits"] = prediction + return outputs + + def predict_headwear_components(self, images, foreground=None, return_presence=False): + from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits + crop = head_crop(images, foreground) + semantic = self._runtime_semantic_features(crop)["raw_spatial"] + prediction=self.headwear_head(crop,semantic,return_presence=return_presence) + if return_presence:return restore_logits(prediction[0],images.shape[-2:]),prediction[1] + return restore_logits(prediction,images.shape[-2:]) + + def predict_ownership(self, images, foreground=None, return_presence=False): + from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits + crop = head_crop(images, foreground) + semantic = self._runtime_semantic_features(crop)["raw_spatial"] + prediction = self.ownership_head(crop, semantic, return_presence=return_presence) + if return_presence:return restore_logits(prediction[0], images.shape[-2:]),prediction[1] + return restore_logits(prediction, images.shape[-2:]) + + def predict_accessories(self, images, foreground=None, return_components=False): + from SkingToolkit.dense_uv_parser.accessories import head_crop, restore_logits + crop = head_crop(images, foreground) + semantic = self._runtime_semantic_features(crop)["raw_spatial"] + logits = self.accessory_head(crop, semantic, return_components=return_components) + if isinstance(logits,tuple):return tuple(restore_logits(x,images.shape[-2:]) for x in logits) + return restore_logits(logits, images.shape[-2:]) + + def predict_projected_outer_uv_occupancy( + self, + atlas_features, + global_context, + ): + if not self.predict_outer_uv_occupancy: + raise ValueError("This parser has no outer UV occupancy head.") + return self.outer_uv_occupancy_head( + atlas_features, + global_context, + ) + + +def count_parameters(model): + return sum(param.numel() for param in model.parameters() if param.requires_grad) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership.py new file mode 100644 index 0000000000000000000000000000000000000000..88497fe6719e774227e639eb0dde8bcc9847fcf4 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership.py @@ -0,0 +1,94 @@ +"""Independent semantic evidence for head layers; existing v101 heads stay frozen.""" +import torch +from torch import nn +from SkingToolkit.dense_uv_parser.accessories import HeadAccessoryHead, connected_object_support + +OWNERSHIP_CLASSES=('abstain','inner_face','inner_hair','eyewear','headphones') + + +class HeadOwnershipHead(HeadAccessoryHead): + def __init__(self,semantic_dim=768,predict_presence=False): + super().__init__(semantic_dim,False) + self.classifier=nn.Conv2d(24,len(OWNERSHIP_CLASSES),1) + nn.init.zeros_(self.classifier.weight) + nn.init.zeros_(self.classifier.bias) + self.presence = nn.Sequential(nn.Linear(5*8*8,128),nn.SiLU(),nn.Dropout(.1),nn.Linear(128,1)) if predict_presence else None + + def predict_presence(self,logits): + if self.presence is None:raise ValueError('No trained headphone presence gate') + evidence=torch.nn.functional.adaptive_avg_pool2d(logits.float().softmax(1),(8,8)).flatten(1) + return self.presence(evidence).flatten() + + def forward(self,crop,semantic_features,return_presence=False): + logits=super().forward(crop,semantic_features) + return (logits,self.predict_presence(logits)) if return_presence else logits + + +def apply_head_ownership(routing,outputs,foreground,renderer,views): + from SkingToolkit.dense_uv_parser.utils import build_static_surface_routing + logits=outputs.get('head_ownership_logits') + if logits is None:return + p=logits.float().softmax(1);confidence,identity=p.max(1) + old=outputs['accessory_logits'].float().softmax(1) + # The new head abstains on hats and outer hair. Preserve those established + # identities; a confident face label can correct geometric false positives. + protected=(old[:,2]+old[:,3])>.5 + outer=torch.zeros_like(foreground);inner=torch.zeros_like(foreground) + # Resolve semantic conflicts at the output texture's resolution. A few + # boundary pixels cannot keep an outer texel whose visible footprint is + # confidently a face. These are learned class votes, independent of RGB. + groups=foreground.shape[0]//len(views) + face_votes=p.new_zeros(groups,4096);counts=p.new_zeros(groups,4096) + for vi,view in enumerate(views): + static=build_static_surface_routing(renderer,view,foreground.device) + valid=foreground[vi::len(views)]&static['masks'][1]&(static['part'][1]==0)&~protected[vi::len(views)] + index=static['flat_uv'][1].flatten()[None].expand(groups,-1) + face_votes.scatter_add_(1,index,(p[vi::len(views),1]*valid).flatten(1)) + counts.scatter_add_(1,index,valid.float().flatten(1)) + face_cells=(face_votes/counts.clamp_min(1)>=.7)&(counts>=4) + veto=torch.zeros_like(foreground) + for vi,view in enumerate(views): + sl=slice(vi,foreground.shape[0],len(views));static=build_static_surface_routing(renderer,view,foreground.device) + valid=foreground[sl]&static['masks'][1]&(static['part'][1]==0) + support,_=connected_object_support(confidence[sl],(identity[sl]==4).long(),valid,.95,.5) + presence=outputs.get('headphone_presence_logit') + if presence is not None: + probabilities=presence.sigmoid() + if outputs.get('headphone_presence_consensus',False): + probabilities=probabilities.reshape(-1,len(views)).amin(1).repeat_interleave(len(views)) + support&=(probabilities[sl]>=outputs.get('headphone_presence_threshold',.9))[:,None,None] + mode=outputs.get('headphone_routing_mode','unrestricted') + if mode=='disabled':support.zero_() + elif mode=='existing_uv': + occupancy=outputs.get('headphone_uv_support') + if occupancy is None:raise ValueError('Headphone colour routing requires accepted UV support') + accepted_uv=occupancy.gather(1,static['flat_uv'][1].flatten()[None].expand(groups,-1)).reshape_as(valid) + support&=accepted_uv + outer[sl]=support&~protected[sl] + cell=face_cells.gather(1,static['flat_uv'][1].flatten()[None].expand(groups,-1)).reshape_as(valid)&valid + veto[sl]=cell + outer[sl]&=~cell + inner[sl]=(cell|((identity[sl]==1)&(confidence[sl]>=.95)&(old[sl,1]<.5)))&~protected[sl]&foreground[sl]&static['masks'][0]&(static['part'][0]==0) + for layer,accepted in ((1,outer[sl]),(0,inner[sl])): + for name,value in [('layer',layer),('foreground',True),('surface',layer),('secondary',False),('secondary_routed',False),('semantic_fallback',False),('consensus_outer_gate_rejected',False)]: + if name in routing:routing[name][sl]=torch.where(accepted,torch.full_like(routing[name][sl],value),routing[name][sl]) + for name in ('flat_uv','part','face','texel_center_score'): + if name in routing:routing[name][sl]=torch.where(accepted,static[name][layer],routing[name][sl]) + routing['confidence'][sl]=torch.where(accepted,confidence[sl],routing['confidence'][sl]) + cm=(2*confidence[sl]-1).clamp_min(0) + for name,value in [('confidence_margin',cm),('confidence_margin_ratio',cm/confidence[sl].clamp_min(1e-6))]: + if name in routing:routing[name][sl]=torch.where(accepted,value,routing[name][sl]) + routing['accessory_supported']=(routing['accessory_supported']|outer)&~inner + presence=outputs.get('headphone_presence_logit') + present=torch.ones(p.shape[0],device=p.device,dtype=torch.bool) + if presence is not None: + probability=presence.sigmoid() + if outputs.get('headphone_presence_consensus',False): + probability=probability.reshape(-1,len(views)).amin(1).repeat_interleave(len(views)) + present=probability>=outputs.get('headphone_presence_threshold',.9) + routing['headphone_colour_conflict']=(p[:,4]>=.5)&present[:,None,None]&foreground + routing['ownership_face_cell_veto']=veto + routing['ownership_outer_supported']=outer + routing['ownership_inner_supported']=inner + routing['head_ownership_identity']=identity + routing['head_ownership_probability']=confidence diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership_data.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership_data.py new file mode 100644 index 0000000000000000000000000000000000000000..6944e806a21a91aeaa3300d02ae3510e89c64476 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/ownership_data.py @@ -0,0 +1,170 @@ +"""Authored head ownership: abstain, inner face, inner hair, eyewear, headphones. + +Real regression examples are never used to produce training labels. +""" +import numpy as np +import torch +from torch.utils.data import Dataset +from SkingToolkit.dense_uv_parser.skin_dataset import load_skin + + +def make_ownership_skin(original, seed): + rng = np.random.default_rng(seed) + uv = original.permute(1, 2, 0).numpy().copy() + objects = np.zeros((64, 64), dtype=np.int64) + ownership = np.zeros((64,64),dtype=np.int64) + components = np.zeros((64,64),dtype=np.int64) # none, inner crown/band, outer crown/band, brim + # All six head faces are authored together. Body pixels remain the source skin. + uv[:16, :32] = 0 + uv[:16, 32:] = 0 + faces = [(8, 8), (24, 8), (0, 8), (16, 8), (8, 0), (16, 0)] + skin = np.array(rng.choice([[.87,.63,.43],[.45,.28,.19],[.65,.44,.30], + [.95,.78,.62],[.32,.20,.13],[.74,.62,.49]])) + if rng.random() < .12: + skin = rng.uniform(.2, .9, 3) + hair = rng.uniform(.03, .30, 3) if rng.random() < .7 else rng.uniform(.1, .85, 3) + for fi, (x, y) in enumerate(faces): + tex = np.clip(skin[None,None] + rng.normal(0,.025,(8,8,1)), 0, 1) + if fi in (1,4): tex[:] = hair + rng.normal(0,.025,(8,8,1)) + if fi in (0,2,3): tex[:rng.integers(1,3)] = hair + uv[y:y+8, x:x+8, :3] = np.clip(tex,0,1) + uv[y:y+8, x:x+8, 3] = 1 + own = ownership[y:y+8,x:x+8] + own[:] = 1 + if fi in (1,4): own[:] = 2 + if fi in (0,2,3): + # Labels follow the authored hair mask, including its sampled depth. + own[np.all(tex == hair,axis=2)] = 2 + # Eyes, eyebrows and mouth are explicit negative facial examples. + for x in (9, 13): + uv[11:12, x:x+2, :3] = rng.uniform(.01,.18,3) + uv[14,11:13,:3] = rng.uniform(.05,.5,3) + if rng.random() < .4: + uv[10,9:11,:3] = hair + uv[10,13:15,:3] = hair + + def put(face, mask, color, identity, layer=1, component=0): + x, y = faces[face]; x += 32*layer + patch = uv[y:y+8, x:x+8] + patch[mask, :3] = color[mask] if np.asarray(color).ndim == 3 else color + patch[mask, 3] = 1 + objects[y:y+8, x:x+8][mask] = identity + components[y:y+8, x:x+8][mask] = component + ownership[y:y+8,x:x+8][mask] = 3 if identity==1 else 0 + + yy, xx = np.mgrid[:8,:8] + negative = rng.random() < .18 + if not negative and rng.random() < .78: + top = int(rng.integers(2,4)); height = int(rng.integers(3,5)) + frame_color = rng.uniform(.025,.23,3) if rng.random()<.70 else rng.uniform(.01,.85,3) + lens_color = rng.uniform(.05,.95,3) + frame = np.zeros((8,8),bool); lenses = np.zeros_like(frame) + # Distinct connected eyewear styles; class ID is independent of palette. + visor = rng.random()<.20 + boxes = [(0,3),(4,7)] if visor else [(0,2),(5,7)] + hollow = rng.random()<.30 + rimless_top = rng.random()<.45 and not hollow + for box_index,(left, right) in enumerate(boxes): + rectangle = (xx>=left)&(xx<=right)&(yy>=top)&(yy.45)) + tex=np.clip(np.broadcast_to(hair,(8,8,3))+rng.normal(0,.025,(8,8,1)),0,1) + put(face,mask,tex,3) + # Headphones are authored as one object across crown and both ears. Palettes + # are independent of object identity; visible underlying hair is labelled inner. + if not (components>0).any() and rng.random()<.55: + band_row=int(rng.integers(1,5));width=int(rng.integers(1,3)) + color=rng.uniform(.02,.9,3); pad_color=rng.uniform(.01,.5,3) + if rng.random()<.25:color=hair.copy() + top=(yy>=band_row)&(yy=column)&(xx=max(0,center-int(rng.integers(1,3))))&(xx<=min(7,center+int(rng.integers(1,3))))&(yy>=3)&(yy<=int(rng.integers(5,8))) + tex=np.broadcast_to(color,(8,8,3)).copy() + tex[pad]=pad_color + accent=pad&(xx>=center)&(yy>=4)&(yy<=5);tex[accent]=color + mask=stem|pad;put(face,mask,tex,0) + x,y=faces[face];ownership[y:y+8,x+32:x+40][mask]=4 + # Side hair underneath an opaque earcup is not an observed earcup. + inner=uv[y:y+8,x:x+8,:3];inner[:,:]=np.clip(hair+rng.normal(0,.025,(8,8,1)),0,1) + ownership[y:y+8,x:x+8]=2 + return torch.from_numpy(uv).permute(2,0,1).float(),torch.from_numpy(ownership) + + +class OwnershipDataset(Dataset): + def __init__(self,paths,count,seed): + self.paths,self.count,self.seed,self.epoch=list(paths),count,seed,0 + def __len__(self):return self.count + def __getitem__(self,index): + seed=self.seed+index*104729+self.epoch*self.count*104729 + uv,labels=make_ownership_skin(load_skin(self.paths[(index+self.epoch*self.count)%len(self.paths)]),seed) + return {'uv':uv,'ownership':labels} + + +def render_ownership(uv,labels,renderer,views): + from SkingToolkit.dense_uv_parser.utils import build_dense_parser_batch + images,targets=[],[] + for view in views: + image,t=build_dense_parser_batch(uv,renderer,view) + xy=(t['uv']*63).round().long().clamp(0,63);index=xy[:,1]*64+xy[:,0] + target=labels.flatten(1).gather(1,index.flatten(1)).reshape_as(index) + target=target.masked_fill((t['foreground'][:,0]<.5)|(t['part']!=0),0) + images.append(image);targets.append(target) + return torch.stack(images,1).flatten(0,1),torch.stack(targets,1).flatten(0,1) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/prepare_foreground_data.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/prepare_foreground_data.py new file mode 100644 index 0000000000000000000000000000000000000000..127d8b593104d68f03cf42695bbc16680ec4817f --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/prepare_foreground_data.py @@ -0,0 +1,80 @@ +"""Render identity-disjoint supervised RGBA cutouts; no flood/pseudo alpha labels.""" +import argparse,hashlib,json,os,shutil,time +from pathlib import Path +import numpy as np +import torch +from PIL import Image +from SkingToolkit.renderer import DifferentiableRenderer +from SkingToolkit.dense_uv_parser.skin_dataset import load_skin +from SkingToolkit.dense_uv_parser.semantic_generalization import write_json + + +def deduplicate_manifest(output): + path=Path(output)/'manifest.json';manifest=json.loads(path.read_text()) + if 'canonical_identity_check' in manifest:return + hashes={source:hashlib.sha256(load_skin(source).numpy().tobytes()).hexdigest() for source in {r['source'] for r in manifest['records']}} + ownership={};kept=[] + for split in ('test','validation','train'): + for record in manifest['records']: + if record['split']!=split:continue + key=hashes[record['source']] + if key not in ownership:ownership[key]=(split,record['source']) + if ownership[key]!=(split,record['source']):continue + kept.append({**record,'canonical_rgba_sha256':key}) + manifest['records']=kept + manifest['counts']={split:len({r['canonical_rgba_sha256'] for r in kept if r['split']==split}) for split in ('train','validation','test')} + manifest['canonical_identity_check']='Canonical normalized RGBA hashed; no identical skin in multiple splits; test then validation take precedence' + write_json(path,manifest) + print(json.dumps({'deduplicated_sources':manifest['counts']}),flush=True) + + +def main(): + root=Path(__file__).resolve().parents[1] + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--output-dir',type=Path,required=True) + p.add_argument('--train-sources',type=int,default=4096) + p.add_argument('--val-sources',type=int,default=128) + p.add_argument('--test-sources',type=int,default=128) + o=p.parse_args();o.output_dir.mkdir(parents=True,exist_ok=False) + torch.set_num_threads(4) + split_path=root/'dense_uv_parser/runs/dense_uv_parser_v101/source_splits.json' + splits=json.loads(split_path.read_text());rng=np.random.default_rng(610106) + selected={} + for name,n in [('train',o.train_sources),('validation',o.val_sources),('test',o.test_sources)]: + selected[name]=[splits[name][int(i)] for i in rng.choice(len(splits[name]),n,replace=False)] + assert not(set(selected['train'])&set(selected['validation']) or set(selected['train'])&set(selected['test']) or set(selected['validation'])&set(selected['test'])) + mapping=o.output_dir/'mappings';mapping.mkdir() + source_mapping=root.parent/'github/differentiable_minecraft_renderer/mappings_512x1024' + views=['front_left','back_left'] + for v in views:(mapping/(v+'_mapping.pt')).symlink_to(source_mapping/(v+'_mapping.pt')) + renderer=DifferentiableRenderer(str(mapping),bg_color=(0,0,0)).cuda() + records=[];start=time.time() + for name,paths in selected.items(): + dest=o.output_dir/name;dest.mkdir() + for begin in range(0,len(paths),4): + batch_paths=paths[begin:begin+4] + uv=torch.stack([load_skin(f) for f in batch_paths]).cuda() + for view in views: + with torch.no_grad():rgba=renderer.forward_view(uv,view).float().clamp(0,1).cpu() + alpha=rgba[:,3:4] + rgb=(rgba[:,:3]/alpha.clamp_min(1/255)).clamp(0,1) + straight=torch.cat([rgb,alpha],1) + for j,source in enumerate(batch_paths): + source_sha=hashlib.sha256(Path(source).read_bytes()).hexdigest() + filename=f'{source_sha[:20]}_{view}.png' + file=dest/filename + Image.fromarray((straight[j].permute(1,2,0).numpy()*255).round().astype(np.uint8),'RGBA').save(file,compress_level=2) + records.append({'split':name,'source':source,'source_sha256':source_sha,'view':view,'file':str(file.resolve())}) + if begin%128==0: + state={'state':'rendering','split':name,'sources':begin+len(batch_paths),'total_sources':len(paths),'elapsed':time.time()-start} + write_json(o.output_dir/'status.json',state);print(json.dumps(state),flush=True) + manifest={'recipe':'Renderer alpha and foreground RGB; PNG is straight RGBA; identity-disjoint source selection', + 'original_split_sha256':hashlib.sha256(split_path.read_bytes()).hexdigest(),'views':views, + 'mapping_sha256':{v:hashlib.sha256((source_mapping/(v+'_mapping.pt')).read_bytes()).hexdigest() for v in views}, + 'counts':{k:len(v) for k,v in selected.items()},'records':records} + write_json(o.output_dir/'manifest.json',manifest) + deduplicate_manifest(o.output_dir) + write_json(o.output_dir/'status.json',{'state':'complete','images':len(records),'elapsed':time.time()-start}) + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/released_inference.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/released_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..dca1fabc78846e75d6e4ad38e958251cd5c8d0e7 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/released_inference.py @@ -0,0 +1,48 @@ +"""Run a checksum-pinned v101 release, including its trained foreground model.""" +import argparse +import hashlib +import json +import subprocess +import sys +from pathlib import Path + + +def main(): + root=Path(__file__).resolve().parents[1] + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--checkpoint',type=Path,help='Explicit parser checkpoint override') + p.add_argument('--release',type=Path,default=Path(__file__).with_name('v101_release.json')) + p.add_argument('--output-dir',type=Path,default=Path(__file__).parent/'output_history/v101') + p.add_argument('--inputs',type=Path,nargs='+') + p.add_argument('--foreground-probabilities',type=Path,nargs='+') + p.add_argument('--save-tensors',action='store_true') + p.add_argument('--pipeline',type=Path) + opt=p.parse_args() + release=json.loads(opt.release.read_text()) + if not release['passed']:raise ValueError('This release has not passed validation') + checkpoint=opt.checkpoint.resolve() if opt.checkpoint else root/release['checkpoint'] + foreground=release['foreground'] + delta=root/foreground['checkpoint'] + checks=[(delta,foreground['checkpoint_sha256'])] + if not opt.checkpoint:checks.append((checkpoint,release['checkpoint_sha256'])) + for path,expected in checks: + with path.open('rb') as f:actual=hashlib.file_digest(f,'sha256').hexdigest() + if actual!=expected:raise ValueError('Release checkpoint changed: '+str(path)) + pipeline=opt.pipeline.resolve() if opt.pipeline else root/release['pipeline'] + if not opt.pipeline and hashlib.sha256(pipeline.read_bytes()).hexdigest()!=release['pipeline_sha256']: + raise ValueError('Release pipeline changed') + entry=root/'dense_uv_parser/run_local.py' + command=[sys.executable,str(entry)] + if opt.foreground_probabilities: + command+=['batch_accessories','--foreground-probabilities',*[str(x.resolve()) for x in opt.foreground_probabilities]] + else: + command+=['foreground_batch','--foreground-checkpoint',str(delta), + '--foreground-model-dir',foreground['model_dir'], + '--probability-dir',str(root/'dense_uv_parser/cache/released_foreground')] + command+=['--checkpoint',str(checkpoint),'--pipeline',str(pipeline),'--output-dir',str(opt.output_dir.resolve())] + if opt.inputs:command+=['--inputs',*[str(x.resolve()) for x in opt.inputs]] + if opt.save_tensors:command+=['--save-tensors'] + subprocess.run(command,cwd=root,check=True) + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/run_local.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/run_local.py new file mode 100644 index 0000000000000000000000000000000000000000..e93fa13bd80596dfe393725bcfcf5cad70d0f61e --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/run_local.py @@ -0,0 +1,35 @@ +"""Run a toolkit entry point using only this checkout, regardless of its name. + +Example: python dense_uv_parser/run_local.py semantic_generalization --help +""" +import importlib.machinery +import runpy +import sys +import types +from pathlib import Path + + +def bind_checkout(): + root = Path(__file__).resolve().parents[1] + existing = sys.modules.get("SkingToolkit") + if existing is not None: + if list(existing.__path__) != [str(root)]: + raise RuntimeError("SkingToolkit was already imported from another checkout") + else: + package = types.ModuleType("SkingToolkit") + package.__path__ = [str(root)] + package.__package__ = "SkingToolkit" + package.__spec__ = importlib.machinery.ModuleSpec( + "SkingToolkit", loader=None, is_package=True + ) + package.__spec__.submodule_search_locations = package.__path__ + sys.modules["SkingToolkit"] = package + return root + + +if __name__ == "__main__": + bind_checkout() + entry = sys.argv.pop(1) if len(sys.argv) > 1 else "semantic_generalization" + if entry not in {"semantic_generalization", "infer", "train", "test_semantic_generalization", "train_accessories", "test_accessories", "batch_accessories", "test_affine_routing", "test_semantic_parser", "test_foreground", "foreground_provider", "prepare_foreground_data", "train_foreground", "test_matting_data", "foreground_batch", "released_inference", "train_ownership", "test_ownership", "train_headphone_presence", "train_headwear", "test_headwear", "test_crown_geometry", "train_headwear_presence"}: + raise SystemExit("Unsupported local entry point") + runpy.run_module(f"SkingToolkit.dense_uv_parser.{entry}", run_name="__main__", alter_sys=True) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/runtime.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..b7433ba4388ea9d4fef67845372077441faed489 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/runtime.py @@ -0,0 +1,7 @@ +import torch + + +def get_device(device_arg): + if device_arg == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + return torch.device(device_arg) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..6901e77bb6acd62a9a5b5450a929e6e4c92e12dc --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic.py @@ -0,0 +1,145 @@ +"""Runtime and cached frozen semantics for Dense UV Parser.""" + +from pathlib import Path + +import torch + +from SkingToolkit.dense_uv_parser.semantic_backbone import ( + SigLIP2VisionBackbone, + TIPSv2VisionBackbone, +) + + +def build_semantic_runtime( + backbone_name, + model_name, + device, + semantic_channels=128, + local_files_only=False, + runtime_batch_size=32, +): + if backbone_name == "siglip2": + backbone = SigLIP2VisionBackbone( + model_name=model_name, + token_channels=semantic_channels, + local_files_only=local_files_only, + ) + elif backbone_name == "tipsv2": + backbone = TIPSv2VisionBackbone( + model_name=model_name, + local_files_only=local_files_only, + inference_batch_size=runtime_batch_size, + ) + else: + raise ValueError(f"Unsupported semantic backbone: {backbone_name!r}.") + backbone = backbone.to(device).eval() + backbone.requires_grad_(False) + return backbone + + +def attach_semantic_runtime( + model, + backbone_name, + model_name, + device, + local_files_only=False, + backbone=None, + runtime_batch_size=32, +): + """Attach a frozen tower without adding it to parser checkpoints.""" + if ( + model.semantic_feature_dim <= 0 + and getattr(model, "semantic_spatial_feature_dim", 0) <= 0 + ): + return None + if backbone is None: + backbone = build_semantic_runtime( + backbone_name, + model_name, + device, + semantic_channels=model.semantic_channels, + local_files_only=local_files_only, + runtime_batch_size=runtime_batch_size, + ) + if ( + model.semantic_feature_dim > 0 + and backbone.raw_feature_dim != model.semantic_feature_dim + ): + raise ValueError( + "Semantic feature dimension does not match the parser checkpoint: " + f"backbone={backbone.raw_feature_dim}, parser={model.semantic_feature_dim}." + ) + backbone_spatial_dim = int( + getattr(backbone, "raw_spatial_feature_dim", 0) + ) + parser_spatial_dim = int( + getattr(model, "semantic_spatial_feature_dim", 0) + ) + if parser_spatial_dim > 0 and backbone_spatial_dim != parser_spatial_dim: + raise ValueError( + "Semantic spatial feature dimension does not match the parser " + f"checkpoint: backbone={backbone_spatial_dim}, " + f"parser={parser_spatial_dim}." + ) + # Bypass nn.Module registration: the frozen pretrained tower is referenced + # by name in model_config and must not be duplicated in every parser state_dict. + object.__setattr__(model, "_runtime_semantic_backbone", backbone) + return backbone + + +def build_siglip_runtime( + model_name, + device, + semantic_channels=128, + local_files_only=False, + runtime_batch_size=32, +): + """Backward-compatible wrapper for existing SigLIP2 tooling.""" + return build_semantic_runtime( + "siglip2", + model_name, + device, + semantic_channels=semantic_channels, + local_files_only=local_files_only, + runtime_batch_size=runtime_batch_size, + ) + + +def attach_siglip_runtime( + model, + model_name, + device, + local_files_only=False, + backbone=None, + runtime_batch_size=32, +): + """Backward-compatible wrapper for existing SigLIP2 checkpoints.""" + return attach_semantic_runtime( + model, + "siglip2", + model_name, + device, + local_files_only=local_files_only, + backbone=backbone, + runtime_batch_size=runtime_batch_size, + ) + + +def cached_semantic_batch(cache, paths, device): + if cache is None: + return None + filenames = [Path(path).name for path in paths] + global_features = torch.stack( + [cache.get(filename) for filename in filenames], + dim=0, + ).to(device=device, non_blocking=True) + if not getattr(cache, "has_spatial", False): + return global_features + spatial_features = torch.stack( + [cache.get_spatial(filename) for filename in filenames], + dim=0, + ).to(device=device, non_blocking=True) + return { + "raw_global": global_features, + "raw_spatial": spatial_features, + } diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_backbone.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..02ec73f3ea2ce94f573841ecf891dd75cc3f494f --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_backbone.py @@ -0,0 +1,441 @@ +"""Frozen open-vocabulary visual semantics for Dense UV Parser. + +The dependency on ``transformers`` is intentionally lazy. Dataset utilities, +unit tests, and the geometry-only fallback therefore remain usable on machines +without a Hugging Face training environment. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class SigLIP2VisionBackbone(nn.Module): + """Frozen SigLIP2 vision tower with differentiable image preprocessing. + + FixRes SigLIP2 checkpoints are used deliberately: a differentiable, + aspect-preserving letterbox can be shared by source renders and predicted + renders. That lets the render-semantic cycle propagate gradients all the + way back to the predicted UV atlas. + """ + + def __init__( + self, + model_name="google/siglip2-base-patch16-224", + token_channels=128, + local_files_only=False, + ): + super().__init__() + if "naflex" in model_name.lower(): + raise ValueError( + "NaFlex checkpoints are not supported by the differentiable render-cycle " + "adapter. Use a FixRes SigLIP2 checkpoint such as " + "google/siglip2-base-patch16-224." + ) + try: + from transformers import AutoImageProcessor, AutoModel + except ImportError as error: + raise ImportError( + "SigLIP2 training requires Hugging Face Transformers. Install the remote " + "training environment with: pip install -U transformers sentencepiece " + "safetensors" + ) from error + + processor = AutoImageProcessor.from_pretrained( + model_name, + local_files_only=bool(local_files_only), + use_fast=False, + ) + full_model = AutoModel.from_pretrained( + model_name, local_files_only=bool(local_files_only) + ) + if not hasattr(full_model, "vision_model"): + raise ValueError(f"{model_name} does not expose a SigLIP vision tower.") + self.vision_model = full_model.vision_model + config = getattr(full_model.config, "vision_config", self.vision_model.config) + image_size = getattr(config, "image_size", None) + patch_size = getattr(config, "patch_size", None) + hidden_size = getattr(config, "hidden_size", None) + if image_size is None: + processor_size = getattr(processor, "size", {}) + if isinstance(processor_size, dict): + image_size = processor_size.get("height") or processor_size.get("shortest_edge") + if not isinstance(image_size, int) or not isinstance(patch_size, int): + raise ValueError( + f"{model_name} is not a supported fixed-resolution SigLIP2 vision model." + ) + if not isinstance(hidden_size, int): + raise ValueError(f"Cannot determine hidden_size for {model_name}.") + + image_mean = getattr(processor, "image_mean", (0.5, 0.5, 0.5)) + image_std = getattr(processor, "image_std", (0.5, 0.5, 0.5)) + self.model_name = str(model_name) + self.image_size = int(image_size) + self.patch_size = int(patch_size) + self.raw_feature_dim = int(hidden_size) + self.raw_spatial_feature_dim = int(hidden_size) + self.token_channels = int(token_channels) + # Cache only immutable Python index tuples. A CUDA tensor kept here can + # outlive a torch.compile CUDA Graph invocation and point at storage + # overwritten by the next replay. + self._valid_token_index_cache = {} + self.register_buffer( + "image_mean", + torch.tensor(image_mean, dtype=torch.float32).view(1, 3, 1, 1), + persistent=False, + ) + self.register_buffer( + "image_std", + torch.tensor(image_std, dtype=torch.float32).view(1, 3, 1, 1), + persistent=False, + ) + self.token_projection = nn.Sequential( + nn.LayerNorm(hidden_size), + nn.Linear(hidden_size, token_channels), + ) + self.global_projection = nn.Sequential( + nn.LayerNorm(hidden_size), + nn.Linear(hidden_size, token_channels), + ) + + self.vision_model.requires_grad_(False) + self.vision_model.eval() + + def train(self, mode=True): + super().train(mode) + # The projection adapters train; the pretrained vision tower never does. + self.vision_model.eval() + return self + + def project_global(self, raw_global): + """Project cached or freshly encoded language-aligned global features.""" + if raw_global.shape[-1] != self.raw_feature_dim: + raise ValueError( + f"Expected SigLIP feature dim {self.raw_feature_dim}, " + f"got {raw_global.shape[-1]}." + ) + return self.global_projection(raw_global) + + def _letterbox(self, images): + if images.dim() != 4 or images.shape[1] != 3: + raise ValueError(f"Expected Nx3xHxW images, got {tuple(images.shape)}.") + height, width = images.shape[-2:] + scale = min(self.image_size / height, self.image_size / width) + resized_height = max(1, min(self.image_size, round(height * scale))) + resized_width = max(1, min(self.image_size, round(width * scale))) + resized = F.interpolate( + images.float(), + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + pad_height = self.image_size - resized_height + pad_width = self.image_size - resized_width + left = pad_width // 2 + right = pad_width - left + top = pad_height // 2 + bottom = pad_height - top + pixels = F.pad(resized, (left, right, top, bottom), value=0.5) + pixels = (pixels - self.image_mean) / self.image_std + return pixels, (top, left, resized_height, resized_width) + + def _valid_token_indices(self, content_rect, token_count, device): + cache_key = (*content_rect, int(token_count)) + cached = self._valid_token_index_cache.get(cache_key) + if cached is None: + top, left, resized_height, resized_width = content_rect + patch_grid = self.image_size // self.patch_size + patch_tokens = patch_grid * patch_grid + if token_count not in (patch_tokens, patch_tokens + 1): + cached = tuple(range(token_count)) + else: + row_start = top // self.patch_size + row_stop = min( + patch_grid, + (top + resized_height + self.patch_size - 1) // self.patch_size, + ) + column_start = left // self.patch_size + column_stop = min( + patch_grid, + (left + resized_width + self.patch_size - 1) // self.patch_size, + ) + patch_indices = tuple( + row * patch_grid + column + for row in range(row_start, row_stop) + for column in range(column_start, column_stop) + ) + cached = ( + (0, *(index + 1 for index in patch_indices)) + if token_count == patch_tokens + 1 + else patch_indices + ) + self._valid_token_index_cache[cache_key] = cached + # This tensor belongs only to the current eager/compiled invocation. + # Never retain it in Python state across CUDA Graph replays. + return torch.tensor(cached, dtype=torch.long, device=device) + + def encode_global(self, images): + """Encode only pooled semantics, skipping patch projection and masking.""" + pixels, _ = self._letterbox(images) + outputs = self.vision_model(pixel_values=pixels) + raw_global = outputs.pooler_output + return { + "global": self.project_global(raw_global), + "raw_global": raw_global, + } + + def encode_dense(self, images): + """Return pooled and full-width spatial features without trainable projection.""" + pixels, content_rect = self._letterbox(images) + outputs = self.vision_model(pixel_values=pixels) + raw_tokens = outputs.last_hidden_state + patch_grid = self.image_size // self.patch_size + patch_token_count = patch_grid * patch_grid + if raw_tokens.shape[1] == patch_token_count + 1: + raw_tokens = raw_tokens[:, 1:] + elif raw_tokens.shape[1] != patch_token_count: + raise ValueError( + f"Expected {patch_token_count} SigLIP2 patch tokens, got " + f"{raw_tokens.shape[1]}." + ) + raw_spatial = raw_tokens.transpose(1, 2).reshape( + raw_tokens.shape[0], + raw_tokens.shape[2], + patch_grid, + patch_grid, + ) + top, left, resized_height, resized_width = content_rect + row_start = top // self.patch_size + row_stop = min( + patch_grid, + (top + resized_height + self.patch_size - 1) // self.patch_size, + ) + column_start = left // self.patch_size + column_stop = min( + patch_grid, + (left + resized_width + self.patch_size - 1) // self.patch_size, + ) + raw_spatial = raw_spatial[ + :, :, row_start:row_stop, column_start:column_stop + ] + return { + "raw_global": outputs.pooler_output, + "raw_spatial": raw_spatial, + } + + def forward(self, images): + pixels, content_rect = self._letterbox(images) + outputs = self.vision_model( + pixel_values=pixels, + ) + raw_tokens = outputs.last_hidden_state + raw_global = outputs.pooler_output + + # Every image in this call has the same spatial shape and therefore the + # same letterbox mask. Remove padded patch tokens instead of carrying a + # key-padding mask into every UV cross-attention layer. Besides reducing + # memory length, an unmasked attention call can use PyTorch's fused SDPA + # kernels on supported CUDA devices. + valid_token_indices = self._valid_token_indices( + content_rect, raw_tokens.shape[1], raw_tokens.device + ) + raw_tokens = raw_tokens.index_select(1, valid_token_indices) + patch_mask = torch.ones( + raw_tokens.shape[:2], dtype=torch.bool, device=raw_tokens.device + ) + + return { + "tokens": self.token_projection(raw_tokens), + "token_mask": patch_mask, + "tokens_compact": True, + "global": self.project_global(raw_global), + "raw_global": raw_global, + } + + +class TIPSv2VisionBackbone(nn.Module): + """Frozen TIPSv2 vision backbone exposing its native spatial feature map. + + Dense UV routing needs to distinguish nearby surfaces such as an eye from + hair or a hat brim. A pooled image embedding discards exactly that spatial + evidence, so this adapter uses either the native Transformers + ``AutoBackbone`` feature maps or the official model repository's + ``encode_image().patch_tokens`` compatibility interface. + """ + + def __init__( + self, + model_name="google/tipsv2-b14", + local_files_only=False, + inference_batch_size=32, + ): + super().__init__() + try: + from transformers import AutoModel + except ImportError as error: + raise ImportError( + "TIPSv2 training requires Hugging Face Transformers. Install " + "with: pip install -U transformers safetensors" + ) from error + + native_error = None + try: + from transformers import AutoBackbone + + self.vision_model = AutoBackbone.from_pretrained( + model_name, + local_files_only=bool(local_files_only), + out_indices=[-1], + # Do not let this capability probe ask an interactive question. + # Native releases ignore repository code and load normally; + # older releases fail here and use the explicit trusted fallback. + trust_remote_code=False, + ) + self.runtime_kind = "native_backbone" + except (ImportError, AttributeError, KeyError, OSError, TypeError, ValueError) as error: + native_error = error + try: + # Released TIPSv2 repositories expose a lightweight compatibility + # API before every Transformers build recognizes model_type=tipsv2. + self.vision_model = AutoModel.from_pretrained( + model_name, + local_files_only=bool(local_files_only), + trust_remote_code=True, + ) + self.runtime_kind = "remote_encode_image" + except Exception as remote_error: + raise RuntimeError( + "Unable to load TIPSv2 through either native AutoBackbone " + "or the official trust_remote_code compatibility path. " + f"Native error: {native_error}. Remote error: {remote_error}." + ) from remote_error + + config = self.vision_model.config + image_size = getattr(config, "image_size", getattr(config, "img_size", None)) + if isinstance(image_size, (tuple, list)): + if len(image_size) != 2 or image_size[0] != image_size[1]: + raise ValueError( + f"{model_name} must use a square fixed-resolution input." + ) + image_size = image_size[0] + channels = getattr(self.vision_model, "channels", None) + hidden_size = ( + channels[-1] + if channels + else getattr(config, "hidden_size", getattr(config, "embed_dim", None)) + ) + if not isinstance(image_size, int): + raise ValueError(f"Cannot determine input image size for {model_name}.") + if not isinstance(hidden_size, int): + raise ValueError(f"Cannot determine spatial feature size for {model_name}.") + + self.model_name = str(model_name) + self.image_size = int(image_size) + self.raw_feature_dim = int(hidden_size) + self.raw_spatial_feature_dim = int(hidden_size) + self.inference_batch_size = int(inference_batch_size) + if self.inference_batch_size < 1: + raise ValueError("TIPSv2 inference_batch_size must be positive.") + self.vision_model.requires_grad_(False) + self.vision_model.eval() + + def train(self, mode=True): + super().train(mode) + self.vision_model.eval() + return self + + def _letterbox(self, images): + if images.dim() != 4 or images.shape[1] != 3: + raise ValueError(f"Expected Nx3xHxW images, got {tuple(images.shape)}.") + height, width = images.shape[-2:] + scale = min(self.image_size / height, self.image_size / width) + resized_height = max(1, min(self.image_size, round(height * scale))) + resized_width = max(1, min(self.image_size, round(width * scale))) + resized = F.interpolate( + images.float(), + size=(resized_height, resized_width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + pad_height = self.image_size - resized_height + pad_width = self.image_size - resized_width + left = pad_width // 2 + right = pad_width - left + top = pad_height // 2 + bottom = pad_height - top + pixels = F.pad(resized, (left, right, top, bottom), value=0.5) + return pixels, (top, left, resized_height, resized_width) + + def _crop_letterbox_features(self, features, content_rect): + """Remove letterbox-only cells while retaining the source aspect ratio.""" + top, left, resized_height, resized_width = content_rect + feature_height, feature_width = features.shape[-2:] + row_start = int(math.floor(top * feature_height / self.image_size)) + row_stop = int( + math.ceil((top + resized_height) * feature_height / self.image_size) + ) + column_start = int(math.floor(left * feature_width / self.image_size)) + column_stop = int( + math.ceil((left + resized_width) * feature_width / self.image_size) + ) + row_start = max(0, min(feature_height - 1, row_start)) + row_stop = max(row_start + 1, min(feature_height, row_stop)) + column_start = max(0, min(feature_width - 1, column_start)) + column_stop = max(column_start + 1, min(feature_width, column_stop)) + return features[:, :, row_start:row_stop, column_start:column_stop] + + def encode_dense(self, images): + pixels, content_rect = self._letterbox(images) + spatial_chunks = [] + global_chunks = [] + for start in range(0, pixels.shape[0], self.inference_batch_size): + pixel_chunk = pixels[start : start + self.inference_batch_size] + if self.runtime_kind == "native_backbone": + outputs = self.vision_model(pixel_values=pixel_chunk) + feature_maps = getattr(outputs, "feature_maps", None) + if not feature_maps: + raise ValueError( + f"{self.model_name} did not return spatial backbone feature maps." + ) + spatial_chunk = feature_maps[-1] + global_chunk = spatial_chunk.mean(dim=(2, 3)) + else: + outputs = self.vision_model.encode_image(pixel_chunk) + patch_tokens = getattr(outputs, "patch_tokens", None) + cls_token = getattr(outputs, "cls_token", None) + if patch_tokens is None or cls_token is None: + raise ValueError( + f"{self.model_name} encode_image output must contain " + "patch_tokens and cls_token." + ) + patch_grid = math.isqrt(patch_tokens.shape[1]) + if patch_grid * patch_grid != patch_tokens.shape[1]: + raise ValueError( + "TIPSv2 patch token count must form a square grid; " + f"got {patch_tokens.shape[1]}." + ) + spatial_chunk = patch_tokens.transpose(1, 2).reshape( + patch_tokens.shape[0], + patch_tokens.shape[2], + patch_grid, + patch_grid, + ) + global_chunk = cls_token[:, 0] + spatial_chunks.append(spatial_chunk) + global_chunks.append(global_chunk) + feature_map = torch.cat(spatial_chunks, dim=0) + raw_global = torch.cat(global_chunks, dim=0) + raw_spatial = self._crop_letterbox_features( + feature_map, content_rect + ) + return { + "raw_global": raw_global, + "raw_spatial": raw_spatial, + } + + def forward(self, images): + return self.encode_dense(images) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_cache.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..e3dab315cd5c65ba92246564c1b53ab8a79b9433 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_cache.py @@ -0,0 +1,135 @@ +"""Memory-mapped frozen semantic features used by Dense UV Parser.""" + +import json +from pathlib import Path + +import numpy as np +import torch + + +SIGLIP_CACHE_VERSION = 2 +SUPPORTED_SIGLIP_CACHE_VERSIONS = (1, 2) + + +class SigLIPGlobalCache: + def __init__( + self, + cache_dir, + expected_views=None, + expected_model=None, + expected_data_dir=None, + require_spatial=False, + ): + self.cache_dir = Path(cache_dir) + metadata_path = self.cache_dir / "metadata.json" + embeddings_path = self.cache_dir / "embeddings.npy" + if not metadata_path.is_file() or not embeddings_path.is_file(): + raise FileNotFoundError( + f"Incomplete SigLIP cache in {self.cache_dir}; expected metadata.json " + "and embeddings.npy." + ) + self.metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + if self.metadata.get("version") not in SUPPORTED_SIGLIP_CACHE_VERSIONS: + raise ValueError( + f"Unsupported SigLIP cache version {self.metadata.get('version')!r}." + ) + if expected_views is not None and self.metadata.get("views") != list( + expected_views + ): + raise ValueError( + f"SigLIP cache views={self.metadata.get('views')!r}, " + f"expected {list(expected_views)!r}." + ) + if ( + expected_model is not None + and self.metadata.get("siglip_model") != expected_model + ): + raise ValueError( + f"SigLIP cache model={self.metadata.get('siglip_model')!r}, " + f"expected {expected_model!r}." + ) + if expected_data_dir is not None: + cached_data_dir = self.metadata.get("data_dir") + requested_data_dir = Path(expected_data_dir).resolve() + if cached_data_dir and Path(cached_data_dir).resolve() != requested_data_dir: + raise ValueError( + f"SigLIP cache data_dir={str(Path(cached_data_dir).resolve())!r}, " + f"expected {str(requested_data_dir)!r}." + ) + + filenames = self.metadata.get("filenames") + if not isinstance(filenames, list) or len(set(filenames)) != len(filenames): + raise ValueError("SigLIP cache filenames must be a unique list.") + self.filename_to_index = { + filename: index for index, filename in enumerate(filenames) + } + self.embeddings = np.load(embeddings_path, mmap_mode="r") + expected_shape = ( + len(filenames), + len(self.metadata["views"]), + int(self.metadata["feature_dim"]), + ) + if self.embeddings.shape != expected_shape: + raise ValueError( + f"SigLIP cache shape={self.embeddings.shape}, expected {expected_shape}." + ) + + spatial_metadata = ( + self.metadata.get("spatial_feature_dim"), + self.metadata.get("spatial_height"), + self.metadata.get("spatial_width"), + ) + self.has_spatial = all(value is not None for value in spatial_metadata) + self.spatial_embeddings = None + if self.has_spatial: + spatial_path = self.cache_dir / "spatial_embeddings.npy" + if not spatial_path.is_file(): + raise FileNotFoundError( + "SigLIP cache metadata declares spatial features but " + f"{spatial_path} is missing." + ) + self.spatial_embeddings = np.load(spatial_path, mmap_mode="r") + expected_spatial_shape = ( + len(filenames), + len(self.metadata["views"]), + int(self.metadata["spatial_feature_dim"]), + int(self.metadata["spatial_height"]), + int(self.metadata["spatial_width"]), + ) + if self.spatial_embeddings.shape != expected_spatial_shape: + raise ValueError( + "SigLIP spatial cache shape=" + f"{self.spatial_embeddings.shape}, expected " + f"{expected_spatial_shape}." + ) + elif any(value is not None for value in spatial_metadata): + raise ValueError("SigLIP spatial cache metadata is incomplete.") + if require_spatial and not self.has_spatial: + raise ValueError( + f"SigLIP cache {self.cache_dir} contains only global features; " + "rebuild it with spatial caching enabled." + ) + + def get(self, filename): + try: + index = self.filename_to_index[filename] + except KeyError as error: + raise KeyError(f"{filename} is missing from the SigLIP cache.") from error + return torch.from_numpy( + np.array(self.embeddings[index], dtype=np.float32, copy=True) + ) + + def get_spatial(self, filename): + if not self.has_spatial: + raise ValueError("This SigLIP cache has no spatial features.") + try: + index = self.filename_to_index[filename] + except KeyError as error: + raise KeyError(f"{filename} is missing from the SigLIP cache.") from error + return torch.from_numpy( + np.array( + self.spatial_embeddings[index], + dtype=np.float16, + copy=True, + ) + ) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_generalization.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_generalization.py new file mode 100644 index 0000000000000000000000000000000000000000..84606c36e262a20f51f2a2325ceef915599a3d2f --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_generalization.py @@ -0,0 +1,451 @@ +"""Checkpoint-compatible v61 fine tuning with pixel supervision and context regularization. + +Run through run_local.py. Real-image generalization is an experimental objective; +synthetic appearance stress tests do not establish real-image accuracy. +""" +import argparse +import hashlib +import inspect +import json +import math +import os +import shutil +import sys +import time +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace + +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader, Subset, random_split + +from SkingToolkit.dense_uv_parser import train as core +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet +from SkingToolkit.dense_uv_parser.losses import DenseUVParserLoss +from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime +from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset +from SkingToolkit.dense_uv_parser.utils import IGNORE_INDEX +from SkingToolkit.dense_uv_parser.semantic_targets import ( + build_part_layer_masks, build_semantic_attribute_targets, +) +from SkingToolkit.renderer import DifferentiableRenderer + + +def write_json(path, value): + path = Path(path) + temp = path.with_suffix(path.suffix + ".tmp") + temp.write_text(json.dumps(value, indent=2, allow_nan=False)) + os.replace(temp, path) + + +@torch.no_grad() +def appearance_augment(images, foreground, strength=1.0, generator=None): + """Change appearance without moving pixels, alpha, background, or GT labels.""" + rgb = images[:, :3].float() + fg = foreground.float() + n, _, h, w = rgb.shape + def rand(*shape): + return torch.rand(shape, device=rgb.device, generator=generator) + def randn(*shape): + return torch.randn(shape, device=rgb.device, generator=generator) + # Normalize the blur by foreground support to prevent background bleeding. + smooth = F.avg_pool2d(rgb * fg, 3, 1, 1) / F.avg_pool2d(fg, 3, 1, 1).clamp_min(1e-6) + mix = rand(n, 1, 1, 1) * (0.55 * strength) + altered = rgb * (1 - mix) + smooth * mix + illumination = F.interpolate( + randn(n, 1, 8, 4), size=(h, w), mode="bilinear", align_corners=False + ) * (0.045 * strength) + contrast = 1 + (rand(n, 1, 1, 1) - 0.5) * (0.30 * strength) + tint = (rand(n, 3, 1, 1) - 0.5) * (0.08 * strength) + altered = (altered - 0.5) * contrast + 0.5 + illumination + tint + altered = altered + randn(n, 3, h, w) * (0.008 * strength) + result = images.clone() + result[:, :3] = torch.where(fg.bool(), altered.clamp(0, 1), rgb) + return result + + +def flat_foreground_mask(images, targets): + rgb = images[:, :3].float() + variance = (F.avg_pool2d(rgb.square(), 3, 1, 1) + - F.avg_pool2d(rgb, 3, 1, 1).square()).clamp_min(0) + interior = F.avg_pool2d(targets["foreground"].float(), 3, 1, 1) > 0.999 + return (variance.amax(1) < (2 / 255) ** 2) & interior[:, 0] + + +def flat_route_loss(logits, labels, flat): + per_pixel = F.cross_entropy(logits.float(), labels, ignore_index=IGNORE_INDEX, reduction="none") + terms = [per_pixel[flat & (labels == role)].mean() + for role in (0, 1) if (flat & (labels == role)).any()] + return torch.stack(terms).mean() if terms else logits.sum() * 0 + + +def head_outer_hard_positive_loss(logits, targets): + """Mine the worst head-outer pixels using actual renderer labels.""" + head_outer = (targets["route_role"] == 1) & (targets["part"] == 0) + if not head_outer.any(): + return logits.sum() * 0 + errors = -F.log_softmax(logits.float(), dim=1)[:, 1][head_outer] + return errors.topk(max(1, math.ceil(errors.numel() * 0.20))).values.mean() + + +def skip_regularizer(model, probability): + """Drop entire shallow skip inputs per image only during training.""" + def hook(module, inputs): + if not model.training or probability == 0: + return inputs + x, skip = inputs + keep = (torch.rand(skip.shape[0], 1, 1, 1, device=skip.device) >= probability) + return x, skip * keep.to(skip.dtype) + return [block.register_forward_pre_hook(hook) for block in (model.up0, model.up1)] + + +@contextmanager +def semantics_disabled(model): + handles = [] + if model.semantic_spatial_fusion is not None: + handles.append(model.semantic_spatial_fusion.register_forward_hook( + lambda module, inputs, output: torch.zeros_like(output))) + if model.semantic_fusion is not None: + handles.append(model.semantic_fusion.register_forward_hook( + lambda module, inputs, output: (torch.zeros_like(output[0]), output[1]))) + try: + yield + finally: + for handle in handles: + handle.remove() + + +def build_split(dataset, seed, val_split): + val_count = int(len(dataset) * val_split) + return random_split(dataset, [len(dataset) - val_count, val_count], + generator=torch.Generator().manual_seed(seed)) + + +def count_routes(output, target, flat): + labels = target["route_role"] + valid = labels != IGNORE_INDEX + pred = output["layer"].argmax(1) + counts = {} + for role, name in ((0, "inner"), (1, "outer"), (2, "secondary")): + gt, selected = labels == role, (pred == role) & valid + counts[f"{name}_tp"] = int((gt & selected).sum()) + counts[f"{name}_fp"] = int((~gt & selected).sum()) + counts[f"{name}_fn"] = int((gt & ~selected).sum()) + inner_flat = flat & (labels == 0) + counts["flat_inner_count"] = int(inner_flat.sum()) + counts["flat_inner_false_outer"] = int((inner_flat & (pred == 1)).sum()) + head = (target["part"] == 0) & valid + gt_head_outer, pred_head_outer = head & (labels == 1), head & (pred == 1) + counts["head_outer_tp"] = int((gt_head_outer & pred_head_outer).sum()) + counts["head_outer_fp"] = int((~gt_head_outer & pred_head_outer).sum()) + counts["head_outer_fn"] = int((gt_head_outer & ~pred_head_outer).sum()) + return counts + + +def summarize_counts(counts): + values = dict(counts) + for name in ("inner", "outer", "secondary", "head_outer"): + tp, fp, fn = [counts[f"{name}_{k}"] for k in ("tp", "fp", "fn")] + values[f"{name}_precision"] = tp / max(1, tp + fp) + values[f"{name}_recall"] = tp / max(1, tp + fn) + values[f"{name}_iou"] = tp / max(1, tp + fp + fn) + values["flat_inner_false_outer_rate"] = ( + counts["flat_inner_false_outer"] / max(1, counts["flat_inner_count"])) + return values + + +@torch.no_grad() +def evaluate(model, renderer, loader, args, device, ablation=True, hard=True): + model.eval() + totals = {name: {} for name in ("clean", "appearance")} + if ablation: + totals["no_semantics"] = {} + hard_totals = {} + rng = torch.Generator(device=device).manual_seed(90451) + for batch in loader: + uv = batch["uv"].to(device) + images, targets, _, ids = core.build_parser_inputs(uv, renderer, args.views, False, args) + flat = flat_foreground_mask(images, targets) + with core.autocast_context(device, "bf16"): + # Features always correspond to the actual image being classified. + clean_features = model._runtime_semantic_features(images, targets["foreground"]) + for mode in totals: + if mode == "appearance": + actual = appearance_augment(images, targets["foreground"], 1.25, rng) + output = model(actual, view_ids=ids, semantic_foreground=targets["foreground"]) + elif mode == "no_semantics": + with semantics_disabled(model): + output = model(images, view_ids=ids, semantic_features=clean_features) + else: + output = model(images, view_ids=ids, semantic_features=clean_features) + counts = count_routes(output, targets, flat) + for key, value in counts.items(): + totals[mode][key] = totals[mode].get(key, 0) + value + if mode == "clean" and hard: + hm = core.hard_uv_conditioning_metrics(images, output, targets, renderer, args.views, args) + for key, value in hm.items(): + if key.startswith("count_"): + hard_totals[key] = hard_totals.get(key, 0.0) + float(value) + results = {k: summarize_counts(v) for k, v in totals.items()} + if hard_totals: + hard_metrics = {} + for name in ("inner", "outer"): + tp, fp, fn = [hard_totals[f"count_hard_{name}_{k}"] for k in ("tp", "fp", "fn")] + for metric, value in (("precision", tp / max(1, tp + fp)), + ("recall", tp / max(1, tp + fn)), + ("iou", tp / max(1, tp + fp + fn))): + hard_metrics[f"{name}_{metric}"] = value + hard_metrics[f"{name}_rgb_mae"] = ( + hard_totals[f"count_hard_{name}_rgb_abs"] / + max(1, hard_totals[f"count_hard_{name}_rgb_values"])) + results["hard_uv"] = hard_metrics + return results + + +def eligible(metrics, baseline, tolerance): + # A robustness gain cannot purchase a large regression in ordinary routing. + for section in ("clean", "hard_uv"): + for name in ("inner_iou", "outer_iou", "outer_precision", "outer_recall"): + if metrics[section][name] < baseline[section][name] - tolerance: + return False + return True + + +def robustness_score(metrics): + appearance = metrics["appearance"] + return (0.40 * appearance["head_outer_recall"] + + 0.25 * appearance["outer_recall"] + + 0.25 * appearance["outer_iou"] + + 0.10 * appearance["inner_iou"]) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + root = Path(__file__).resolve().parents[1] + parser.add_argument("--checkpoint", type=Path, default=root / "dense_uv_parser/runs/dense_uv_parser_v61/best.pt") + parser.add_argument("--output_dir", type=Path, required=True) + parser.add_argument("--epochs", type=int, default=2) + parser.add_argument("--train_samples", type=int, default=0, help="0 uses all original 162000 training skins") + parser.add_argument("--val_samples", type=int, default=512) + parser.add_argument("--test_samples", type=int, default=256, + help="Additional original validation skins used only for the final report") + parser.add_argument("--batch_size", type=int, default=8) + parser.add_argument("--workers", type=int, default=4) + parser.add_argument("--lr", type=float, default=1e-5) + parser.add_argument("--semantic_lr", type=float, default=5e-5) + parser.add_argument("--skip_dropout", type=float, default=0.10) + parser.add_argument("--eval_every", type=int, default=500) + parser.add_argument("--log_every", type=int, default=25) + parser.add_argument("--max_clean_regression", type=float, default=0.015) + parser.add_argument("--evaluate_only", action="store_true") + opt = parser.parse_args() + if min(opt.epochs, opt.val_samples, opt.batch_size, opt.eval_every, opt.log_every) < 1: + parser.error("epochs, val_samples, batch_size, eval_every, log_every must be positive") + if not 0 <= opt.skip_dropout < 1 or min(opt.train_samples, opt.test_samples, opt.workers) < 0: + parser.error("invalid dropout or sample/worker count") + opt.output_dir.mkdir(parents=True, exist_ok=False) + device = torch.device("cuda") + torch.set_num_threads(8) + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + core.seed_everything(1234, reproducible=False) + checkpoint = torch.load(opt.checkpoint, map_location="cpu", weights_only=False) + if checkpoint["model_config"].get("predict_outer_uv_occupancy"): + raise ValueError("Expected the v61 model without an occupancy head") + args_dict = vars(core.build_arg_parser().parse_args(["--data_dir", str(root / "skins")])) + args_dict.update(checkpoint["args"]) + args_dict.update(data_dir=str(root / "skins"), + mappings_dir=str(root.parent / "github/differentiable_minecraft_renderer/mappings_256x512"), + output_dir=str(opt.output_dir), privileged_views="", resume=str(opt.checkpoint), + siglip_local_files_only=True, best_metric="semantic_robustness", + outer_uv_min_source_pixels=33, + lambda_outer_false_negative=1.25, + route_outer_class_weight_cap=1.0) + args = SimpleNamespace(**args_dict) + args.views = core.parse_views(args.views) + args.epochs = checkpoint["epoch"] + opt.epochs + # The original 180k universe and random split MUST precede any pilot subsampling. + dataset = SkinUVDataset(args.data_dir, max_samples=args.max_samples) + original_cache = root / "dense_uv_parser" / checkpoint["args"]["siglip_cache_dir"] / "metadata.json" + if original_cache.is_file(): + original_names = json.loads(original_cache.read_text())["filenames"] + if [path.name for path in dataset.skin_paths] != original_names: + raise ValueError("Dataset ordering differs from the v61 cache; reconstruct the original split first") + train_set, val_set = build_split(dataset, args.seed, args.val_split) + if opt.train_samples: + train_set = Subset(dataset, train_set.indices[:opt.train_samples]) + test_set = Subset(dataset, val_set.indices[opt.val_samples:opt.val_samples + opt.test_samples]) + val_set = Subset(dataset, val_set.indices[:opt.val_samples]) + if not len(train_set) or not len(val_set): + raise ValueError("Empty training/validation split") + if set(train_set.indices) & set(val_set.indices): + raise RuntimeError("Training/validation leakage") + loader_args = dict(batch_size=opt.batch_size, num_workers=opt.workers, pin_memory=True) + if opt.workers: + loader_args.update(persistent_workers=True, prefetch_factor=2) + train_loader = DataLoader(train_set, shuffle=True, + generator=torch.Generator().manual_seed(args.seed + 1), **loader_args) + val_loader = DataLoader(val_set, shuffle=False, **loader_args) + test_loader = DataLoader(test_set, shuffle=False, **loader_args) if len(test_set) else None + model_kwargs = {key: val for key, val in checkpoint["model_config"].items() + if key in inspect.signature(DenseUVParserNet).parameters} + model = DenseUVParserNet(**model_kwargs).to(device) + model.load_state_dict(checkpoint["model"], strict=True) + # Static coordinate bias is frozen: only image-conditioned behavior is adapted. + model.route_role_prior.requires_grad_(False) + attach_semantic_runtime(model, "siglip2", args.siglip_model, device, local_files_only=True) + renderer = DifferentiableRenderer(args.mappings_dir).to(device) + loss_kwargs = {key: value for key, value in args_dict.items() + if key in inspect.signature(DenseUVParserLoss).parameters} + loss_kwargs.update(use_uv=False, affine_translation_limit=model.affine_translation_limit, + affine_log_scale_limit=model.affine_log_scale_limit) + criterion = DenseUVParserLoss(**loss_kwargs).to(device) + masks = tuple(mask.to(device) for mask in build_part_layer_masks()) + semantic, visual = [], [] + for name, parameter in model.named_parameters(): + if parameter.requires_grad: + (semantic if name.startswith(("semantic_", "outer_presence", "outer_coverage")) + else visual).append(parameter) + optimizer = torch.optim.AdamW([ + {"params": visual, "lr": opt.lr, "initial_lr": opt.lr}, + {"params": semantic, "lr": opt.semantic_lr, "initial_lr": opt.semantic_lr}, + ], weight_decay=args.weight_decay) + handles = skip_regularizer(model, opt.skip_dropout) + manifest = {"checkpoint": str(opt.checkpoint.resolve()), + "checkpoint_sha256": hashlib.sha256(opt.checkpoint.read_bytes()).hexdigest(), + "source_epoch": checkpoint["epoch"], "checkout": str(root), + "code_files": {Path(m.__file__).name: str(Path(m.__file__).resolve()) + for m in (core, inspect.getmodule(DenseUVParserNet), inspect.getmodule(DenseUVParserLoss))}, + "options": {k: str(v) if isinstance(v, Path) else v for k, v in vars(opt).items()}, + "train_indices": list(train_set.indices), "val_indices": list(val_set.indices), + "test_indices": list(test_set.indices), + "dataset_filenames_sha256": hashlib.sha256( + "\n".join(p.name for p in dataset.skin_paths).encode()).hexdigest(), + "label_source": "original skin RGBA and v61 renderer surface visibility; no v94 pseudo-labels", + "semantic_features": "online from actual input, frozen SigLIP2", + "optimizer": "fresh AdamW for changed fine-tuning objective; v61 model loaded strictly", + "routing": "v61 thresholds held fixed for hard UV validation, source pixel minimum 33"} + for path in manifest["code_files"].values(): + if not Path(path).is_relative_to(root): + raise RuntimeError(f"Imported outside checkout: {path}") + write_json(opt.output_dir / "config.json", {"args": vars(args), "experiment": manifest}) + source_dir = opt.output_dir / "source" + source_dir.mkdir() + for source in (root / "dense_uv_parser").glob("*.py"): + shutil.copy2(source, source_dir / source.name) + shutil.copy2(root / "renderer.py", source_dir / "renderer.py") + started = time.time() + write_json(opt.output_dir / "status.json", {"state": "baseline", "pid": os.getpid()}) + print("Evaluating untouched v61 baseline", flush=True) + baseline = evaluate(model, renderer, val_loader, args, device) + write_json(opt.output_dir / "baseline.json", baseline) + print("baseline=" + json.dumps(baseline), flush=True) + if opt.evaluate_only: + write_json(opt.output_dir / "status.json", {"state": "evaluated", "pid": os.getpid()}) + return + best_score = robustness_score(baseline) + step = 0 + stop = False + rejection_streak = 0 + args.semantic_generalization = manifest["options"] + for epoch in range(1, opt.epochs + 1): + running = 0.0 + for batch_index, batch in enumerate(train_loader): + model.train() + progress = (step / max(1, opt.epochs * len(train_loader) - 1)) + factor = 0.2 + 0.8 * (1 + math.cos(math.pi * progress)) / 2 + for group in optimizer.param_groups: + group["lr"] = group["initial_lr"] * factor + uv = batch["uv"].to(device, non_blocking=True) + images, targets, _, ids = core.build_parser_inputs(uv, renderer, args.views, True, args) + flat = flat_foreground_mask(images, targets) + augmented = step % 2 == 1 + actual = appearance_augment(images, targets["foreground"]) if augmented else images + optimizer.zero_grad(set_to_none=True) + with core.autocast_context(device, "bf16"): + output = model(actual, view_ids=ids, semantic_foreground=targets["foreground"]) + losses = criterion(output, targets) + loss = losses["loss_total"] + loss = loss + 0.10 * flat_route_loss(output["layer"], targets["route_role"], flat) + loss = loss + 0.15 * head_outer_hard_positive_loss(output["layer"], targets) + attributes = build_semantic_attribute_targets(uv, *masks) + loss = loss + 0.25 * F.binary_cross_entropy_with_logits( + output["outer_presence_logits"], attributes["outer_presence"]) + loss = loss + 0.25 * F.smooth_l1_loss(output["outer_coverage"], attributes["outer_coverage"]) + if not augmented: + geometry = core.differentiable_geometry_losses(images, uv, output, targets, renderer, args.views) + for key in ("soft_uv_rgb", "soft_uv_alpha", "soft_uv_inner_recall", + "soft_uv_outer_recall", "render_rgb", "render_alpha"): + loss = loss + getattr(args, "lambda_" + key) * geometry["loss_" + key] + if not torch.isfinite(loss): + raise FloatingPointError(f"Nonfinite loss at step {step}") + loss.backward() + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True) + optimizer.step() + step += 1 + running += float(loss.detach()) + if step % opt.log_every == 0 or step == 1: + status = {"state": "training", "pid": os.getpid(), "step": step, "epoch": epoch, + "total_steps": opt.epochs * len(train_loader), "loss": running / (batch_index + 1), + "gradient_norm": float(norm), "semantic_lr": optimizer.param_groups[1]["lr"], + "elapsed_seconds": time.time() - started} + write_json(opt.output_dir / "status.json", status) + print(json.dumps(status), flush=True) + if step % opt.eval_every == 0 or batch_index + 1 == len(train_loader): + metrics = evaluate(model, renderer, val_loader, args, device) + score = robustness_score(metrics) + accepted = eligible(metrics, baseline, opt.max_clean_regression) + promoted = accepted and score > best_score + record = {"step": step, "epoch": epoch, "score": score, "clean_guard_pass": accepted, + "promoted": promoted, "metrics": metrics, "elapsed_seconds": time.time() - started} + with (opt.output_dir / "metrics.jsonl").open("a") as stream: + stream.write(json.dumps(record) + "\n") + print("evaluation=" + json.dumps(record), flush=True) + if promoted: + best_score = score + args.finetune_step = step + saved_metrics = {"val": metrics, "semantic_robustness": score} + core.save_checkpoint(opt.output_dir / "latest.pt", model, optimizer, None, + checkpoint["epoch"] + epoch, args, saved_metrics, best_metric=-best_score) + if promoted: + core.save_checkpoint(opt.output_dir / "best.pt", model, optimizer, None, + checkpoint["epoch"] + epoch, args, saved_metrics, best_metric=-best_score) + write_json(opt.output_dir / "last_evaluation.json", record) + rejection_streak = 0 if accepted else rejection_streak + 1 + if rejection_streak >= 2: + stop = True + break + if stop: + break + for handle in handles: + handle.remove() + if test_loader is not None and (opt.output_dir / "best.pt").is_file(): + write_json(opt.output_dir / "status.json", {"state": "final_test", "pid": os.getpid(), "step": step}) + best = torch.load(opt.output_dir / "best.pt", map_location="cpu", weights_only=False) + model.load_state_dict(best["model"], strict=True) + candidate_test = evaluate(model, renderer, test_loader, args, device) + model.load_state_dict(checkpoint["model"], strict=True) + baseline_test = evaluate(model, renderer, test_loader, args, device) + write_json(opt.output_dir / "heldout_test.json", { + "note": "Not used for checkpoint selection; appearance stress is synthetic, not real-image ground truth", + "samples": len(test_set), "candidate": candidate_test, "baseline": baseline_test}) + write_json(opt.output_dir / "status.json", {"state": "stopped_clean_regression" if stop else "complete", + "pid": os.getpid(), "step": step, "best_score": best_score, + "baseline_score": robustness_score(baseline), "best_checkpoint_exists": (opt.output_dir / "best.pt").exists(), + "elapsed_seconds": time.time() - started}) + + +if __name__ == "__main__": + try: + main() + except Exception as error: + if "--output_dir" in sys.argv: + index = sys.argv.index("--output_dir") + 1 + if index < len(sys.argv): + failed_dir = Path(sys.argv[index]) + if failed_dir.is_dir() and not isinstance(error, FileExistsError): + write_json(failed_dir / "status.json", { + "state": "failed", "pid": os.getpid(), + "error": f"{type(error).__name__}: {error}"}) + raise diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_targets.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_targets.py new file mode 100644 index 0000000000000000000000000000000000000000..6d74ecfbac069eeb87fc2976323ff31a1126e4d7 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/semantic_targets.py @@ -0,0 +1,33 @@ +"""Atlas-level semantic targets used by Dense UV Parser.""" + +import torch + +from SkingToolkit.dense_uv_parser.uv_layout import build_part_layer_masks + + +def build_semantic_attribute_targets(target_uv, inner_part_masks, outer_part_masks): + alpha = target_uv[:, 3:4].float() + rgb = target_uv[:, :3].float() + outer_masks = outer_part_masks.unsqueeze(0).to(device=target_uv.device) + outer_weight = alpha.unsqueeze(1) * outer_masks + outer_area = outer_masks.sum(dim=(2, 3, 4)).clamp_min(1.0) + outer_coverage = outer_weight.sum(dim=(2, 3, 4)) / outer_area + outer_presence = (outer_coverage > 0.0).float() + + all_masks = torch.cat([inner_part_masks, outer_part_masks], dim=0) + all_masks = all_masks.unsqueeze(0).to(device=target_uv.device) + color_weight = alpha.unsqueeze(1) * all_masks + color_denominator = color_weight.sum(dim=(3, 4)).clamp_min(1.0) + part_colors = ( + rgb.unsqueeze(1) * color_weight + ).sum(dim=(3, 4)) / color_denominator + color_known = color_weight.sum(dim=(2, 3, 4)) > 0.0 + return { + "outer_presence": outer_presence, + "outer_coverage": outer_coverage, + "part_colors": part_colors, + "part_colors_known": color_known, + } + + +__all__ = ["build_part_layer_masks", "build_semantic_attribute_targets"] diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/simple_inpainting.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/simple_inpainting.py new file mode 100644 index 0000000000000000000000000000000000000000..d4f107330448369209e83e06dd4c4bc771db1826 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/simple_inpainting.py @@ -0,0 +1,143 @@ +"""Deterministic, topology-aware repair of missing inner-layer texels.""" + +import torch + +from SkingToolkit.dense_uv_parser.uv_layout import UV_SIZE +from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology + + +def _nearest_defined_source( + target_index, + defined, + valid, + part, + face, + local_v, + positions, + prefer_same_row=False, +): + source_mask = defined & valid & (part == part[target_index]) + used_same_row = False + if prefer_same_row: + same_row_mask = ( + source_mask + & (face < 4) + & torch.isclose( + local_v, + local_v[target_index], + rtol=0.0, + atol=1e-6, + ) + ) + source_indices = same_row_mask.nonzero(as_tuple=False).flatten() + if source_indices.numel() > 0: + used_same_row = True + else: + source_indices = source_mask.nonzero(as_tuple=False).flatten() + else: + source_indices = source_mask.nonzero(as_tuple=False).flatten() + if source_indices.numel() == 0: + return None, False + squared_distance = ( + positions[source_indices] - positions[target_index] + ).square().sum(dim=1) + return source_indices[squared_distance.argmin()], used_same_row + + +def simple_symmetry_nearest_inpaint(uv, alpha_threshold=0.5): + """Fill unknown inner texels while preserving every outer-layer texel.""" + squeeze_batch = uv.dim() == 3 + if squeeze_batch: + uv = uv.unsqueeze(0) + if uv.dim() != 4 or uv.shape[1:] != (4, UV_SIZE, UV_SIZE): + raise ValueError( + f"Expected 4x{UV_SIZE}x{UV_SIZE} or Bx4x{UV_SIZE}x{UV_SIZE} UV, " + f"got {tuple(uv.shape)}." + ) + if not 0.0 <= alpha_threshold <= 1.0: + raise ValueError("alpha_threshold must be in [0, 1].") + + topology = build_simple_uv_topology() + topology_face = topology.face.reshape(-1) + device = uv.device + valid = topology.valid.reshape(-1).to(device=device) + layer = topology.layer.reshape(-1).to(device=device) + part = topology.part.reshape(-1).to(device=device) + face = topology.face.reshape(-1).to(device=device) + local_v = topology.local_uv.reshape(-1, 2)[:, 1].to( + device=device, + dtype=torch.float32, + ) + mirrored = topology.mirrored_texel.reshape(-1).to(device=device) + positions = topology.world_position.reshape(-1, 3).to( + device=device, + dtype=torch.float32, + ) + result = uv.flatten(2).transpose(1, 2).clone() + stats = [] + + for batch_index in range(result.shape[0]): + original_defined = valid & ( + result[batch_index, :, 3] > float(alpha_threshold) + ) + defined = original_defined.clone() + symmetry_filled = 0 + nearest_filled = 0 + same_row_nearest_filled = 0 + for target_index in topology.inner_fill_order.tolist(): + if bool(defined[target_index]): + continue + mirror_index = int(mirrored[target_index]) + if bool(defined[mirror_index]): + result[batch_index, target_index] = result[ + batch_index, mirror_index + ] + defined[target_index] = True + symmetry_filled += 1 + continue + + source_index, used_same_row = _nearest_defined_source( + target_index, + defined, + valid & (layer == 0), + part, + face, + local_v, + positions, + prefer_same_row=int(topology_face[target_index]) in (2, 3), + ) + if source_index is None: + continue + result[batch_index, target_index] = result[ + batch_index, source_index + ] + defined[target_index] = True + nearest_filled += 1 + same_row_nearest_filled += int(used_same_row) + + unresolved_inner = valid & (layer == 0) & ~defined + stats.append( + { + "known_texels": int(original_defined.sum().item()), + "known_inner_texels": int( + (original_defined & (layer == 0)).sum().item() + ), + "known_outer_texels": int( + (original_defined & (layer == 1)).sum().item() + ), + "symmetry_filled_texels": symmetry_filled, + "nearest_3d_filled_texels": nearest_filled, + "same_row_nearest_filled_texels": same_row_nearest_filled, + "preserved_outer_texels": int((valid & (layer == 1)).sum().item()), + "fill_order": "front_back_rings_side_edges_top_bottom_rings", + "color_sources": "currently_defined_only", + "side_nearest_policy": "same_vertical_row_then_same_part_3d", + "unresolved_texels": int(unresolved_inner.sum().item()), + } + ) + + result[:, ~valid] = 0.0 + result = result.transpose(1, 2).reshape_as(uv).clamp(0.0, 1.0) + if squeeze_batch: + return result[0], stats[0] + return result, stats diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/skin_dataset.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/skin_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..ff208277cbcd022a96c0dfa0ec880837a2a219ec --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/skin_dataset.py @@ -0,0 +1,74 @@ +"""Minimal 64x64 Minecraft skin dataset for Dense UV Parser training.""" + +from pathlib import Path + +import numpy as np +import torch +from PIL import Image +from torch.utils.data import Dataset + +from mc_skin_utils.alice_to_steve import alice_to_steve + + +IMAGE_EXTENSIONS = (".png", ".webp", ".jpg", ".jpeg") + + +def load_skin(path, bg_color=(128, 128, 128), normalize_model=True): + skin = Image.open(path).convert("RGBA") + if skin.size != (64, 64): + raise ValueError(f"Expected a 64x64 Minecraft skin at {path}, got {skin.size}.") + if normalize_model and skin.getpixel((47, 52))[3] == 0: + skin = alice_to_steve(skin) + + skin_np = np.array(skin, dtype=np.uint8) + alpha = skin_np[..., 3] + skin_np[(alpha > 0) & (alpha < 255), 3] = 255 + rgba = torch.from_numpy(skin_np.astype(np.float32) / 255.0).permute(2, 0, 1) + alpha = rgba[3:4] + background = torch.tensor(bg_color, dtype=rgba.dtype).view(3, 1, 1) / 255.0 + rgba[:3] = torch.where(alpha > 0, rgba[:3], background) + return rgba.clamp(0.0, 1.0) + + +class SkinUVDataset(Dataset): + def __init__( + self, + data_dir, + mappings_dir=None, + views=None, + max_samples=None, + bg_color=(128, 128, 128), + normalize_model=True, + **_ignored, + ): + # mappings_dir/views remain accepted so old training invocations keep a + # stable constructor while rendering stays in dense_uv_parser/train.py. + del mappings_dir, views + self.data_dir = Path(data_dir) + self.bg_color = bg_color + self.normalize_model = bool(normalize_model) + self.skin_paths = sorted( + path + for path in self.data_dir.iterdir() + if path.is_file() + and path.suffix.lower() in IMAGE_EXTENSIONS + and not path.name.startswith("half_") + ) + if max_samples is not None: + self.skin_paths = self.skin_paths[:max_samples] + if not self.skin_paths: + raise ValueError(f"No skin images found in {self.data_dir}") + + def __len__(self): + return len(self.skin_paths) + + def __getitem__(self, index): + skin_path = self.skin_paths[index] + return { + "uv": load_skin( + skin_path, + bg_color=self.bg_color, + normalize_model=self.normalize_model, + ), + "path": str(skin_path), + } diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train.py new file mode 100644 index 0000000000000000000000000000000000000000..d054ef1709774fa041a742da00f1225b49f422ef --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train.py @@ -0,0 +1,4247 @@ +import argparse +import json +import math +import os +import random +import sys +from contextlib import nullcontext +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader, random_split +from torchvision.utils import save_image + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +WORKSPACE_ROOT = TOOLKIT_ROOT.parent +if str(WORKSPACE_ROOT) not in sys.path: + sys.path.insert(0, str(WORKSPACE_ROOT)) + +from SkingToolkit.dense_uv_parser.losses import DenseUVParserLoss # noqa: E402 +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet, count_parameters # noqa: E402 +from SkingToolkit.dense_uv_parser.semantic import ( # noqa: E402 + attach_semantic_runtime, + build_semantic_runtime, + cached_semantic_batch, +) +from SkingToolkit.dense_uv_parser.utils import ( # noqa: E402 + FACE_PALETTE, + IGNORE_INDEX, + LAYER_FACE_PALETTE, + LAYER_PALETTE, + PART_PALETTE, + ROUTE_ROLE_PALETTE, + ROUTE_OUTER_PRIMARY, + SPLAT_COLOR_AGGREGATIONS, + UV_SIZE, + aggregate_direct_outer_values_by_view, + attach_projected_outer_uv_occupancy, + build_dense_parser_batch, + build_geometry_grid_debug, + build_static_surface_routing, + canonicalize_dense_targets, + canonicalize_parser_outputs, + colorize_foreground, + colorize_labels, + colorize_surface, + colorize_uv, + combine_layer_face, + flat_uv_to_uv01, + fill_geometry_grid_debug, + overlay_geometry_grid_debug, + parse_views, + prediction_uv01, + randomize_render_background, + render_direct_uv, + render_outer_alpha_uv, + soft_splat_geometry_predictions_to_uv, + splat_deterministic_targets_to_uv_conditioning, + splat_parser_predictions_to_uv_conditioning, + splat_predictions_to_uv_conditioning, + splat_targets_to_uv_conditioning, + surface_class_count, +) +from SkingToolkit.dense_uv_parser.semantic_cache import ( # noqa: E402 + SigLIPGlobalCache, +) +from SkingToolkit.dense_uv_parser.semantic_targets import ( # noqa: E402 + build_part_layer_masks, + build_semantic_attribute_targets, +) +from SkingToolkit.dense_uv_parser.runtime import get_device # noqa: E402 +from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset # noqa: E402 +from SkingToolkit.dense_uv_parser.uv_layout import build_uv_masks # noqa: E402 +from SkingToolkit.dense_uv_parser.uv_topology import ( # noqa: E402 + build_outer_uv_graph, +) +from SkingToolkit.renderer import DifferentiableRenderer # noqa: E402 + +try: + from tqdm import tqdm +except ImportError: + tqdm = None + + +def autocast_context(device, precision): + if precision == "no" or device.type == "cpu": + return nullcontext() + dtype = torch.float16 if precision == "fp16" else torch.bfloat16 + return torch.autocast(device_type=device.type, dtype=dtype) + + +def build_grad_scaler(device, precision): + if precision != "fp16" or device.type != "cuda": + return None + return torch.amp.GradScaler("cuda") + + +def seed_everything(seed, reproducible=True): + """Seed every RNG used by model initialization and the training loop.""" + if seed < 0: + raise ValueError("--seed must be non-negative.") + os.environ["PYTHONHASHSEED"] = str(seed) + if reproducible: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + random.seed(seed) + np.random.seed(seed % (2**32)) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def seed_data_worker(_worker_id): + """Derive Python and NumPy worker RNGs from DataLoader's torch seed.""" + worker_seed = torch.initial_seed() % (2**32) + random.seed(worker_seed) + np.random.seed(worker_seed) + torch.manual_seed(worker_seed) + + +def capture_reproducibility_state(train_generator=None, val_generator=None): + numpy_state = np.random.get_state() + state = { + "python": random.getstate(), + "numpy": { + "bit_generator": numpy_state[0], + "state": numpy_state[1].tolist(), + "position": numpy_state[2], + "has_gauss": numpy_state[3], + "cached_gaussian": numpy_state[4], + }, + "torch": torch.get_rng_state(), + "cuda": ( + torch.cuda.get_rng_state_all() + if torch.cuda.is_available() + else None + ), + } + if train_generator is not None: + state["train_loader_generator"] = train_generator.get_state() + if val_generator is not None: + state["val_loader_generator"] = val_generator.get_state() + return state + + +def restore_reproducibility_state( + state, + train_generator=None, + val_generator=None, +): + if not state: + return False + random.setstate(state["python"]) + numpy_state = state["numpy"] + np.random.set_state( + ( + numpy_state["bit_generator"], + np.asarray(numpy_state["state"], dtype=np.uint32), + int(numpy_state["position"]), + int(numpy_state["has_gauss"]), + float(numpy_state["cached_gaussian"]), + ) + ) + torch.set_rng_state(state["torch"].cpu()) + cuda_states = state.get("cuda") + if cuda_states is not None and torch.cuda.is_available(): + torch.cuda.set_rng_state_all( + [cuda_state.cpu() for cuda_state in cuda_states] + ) + if train_generator is not None and "train_loader_generator" in state: + train_generator.set_state(state["train_loader_generator"].cpu()) + if val_generator is not None and "val_loader_generator" in state: + val_generator.set_state(state["val_loader_generator"].cpu()) + return True + + +def configure_torch(args, device): + if hasattr(torch, "set_float32_matmul_precision"): + torch.set_float32_matmul_precision(args.matmul_precision) + torch.use_deterministic_algorithms( + args.reproducible, + warn_only=args.reproducible and not args.strict_determinism, + ) + if device.type == "cuda": + if args.reproducible: + # Flash, memory-efficient, and cuDNN SDPA can select + # non-deterministic backward kernels. The math implementation is + # slower but follows the configured RNG and deterministic policy. + if hasattr(torch.backends.cuda, "enable_flash_sdp"): + torch.backends.cuda.enable_flash_sdp(False) + if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"): + torch.backends.cuda.enable_mem_efficient_sdp(False) + if hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + if hasattr(torch.backends.cuda, "enable_math_sdp"): + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cudnn.deterministic = args.reproducible + torch.backends.cudnn.benchmark = ( + args.cudnn_benchmark and not args.reproducible + ) + + +def learning_rate_for_epoch(base_lr, epoch, epochs, schedule="cosine", min_lr_ratio=0.05): + """Return an absolute-epoch LR so old checkpoints can resume without scheduler state.""" + if base_lr <= 0: + raise ValueError("base_lr must be positive.") + if not 0.0 <= min_lr_ratio <= 1.0: + raise ValueError("min_lr_ratio must be in [0, 1].") + if schedule == "constant" or epochs <= 1: + return float(base_lr) + if schedule != "cosine": + raise ValueError(f"Unsupported learning-rate schedule {schedule!r}.") + progress = min(max((epoch - 1) / max(epochs - 1, 1), 0.0), 1.0) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return float(base_lr) * (min_lr_ratio + (1.0 - min_lr_ratio) * cosine) + + +def format_metrics( + metric_sums, + count, + inner_recall_weight=0.5, + outer_precision_weight=1.0, + outer_recall_weight=1.0, + outer_iou_weight=0.5, + hard_rgb_selection_weight=1.0, + outer_projection_fp_selection_weight=0.25, + outer_projection_area_selection_weight=0.10, +): + result = {} + for name, value in metric_sums.items(): + denominator = 1 if name.startswith("count_") else max(count, 1) + averaged = value / denominator + result[name] = ( + float(averaged.detach().cpu()) + if torch.is_tensor(averaged) + else averaged + ) + + for family in ("", "hard_"): + for role_name in ("inner", "outer"): + count_names = tuple( + f"count_{family}{role_name}_{suffix}" + for suffix in ("tp", "fp", "fn") + ) + if all(name in result for name in count_names): + role_tp = result[f"count_{family}{role_name}_tp"] + role_fp = result[f"count_{family}{role_name}_fp"] + role_fn = result[f"count_{family}{role_name}_fn"] + result[f"{family}precision_{role_name}"] = role_tp / max( + role_tp + role_fp, 1.0 + ) + result[f"{family}recall_{role_name}"] = role_tp / max( + role_tp + role_fn, 1.0 + ) + result[f"{family}iou_{role_name}"] = role_tp / max( + role_tp + role_fp + role_fn, 1.0 + ) + rgb_abs_name = f"count_{family}{role_name}_rgb_abs" + rgb_values_name = f"count_{family}{role_name}_rgb_values" + if rgb_abs_name in result and rgb_values_name in result: + result[f"{family}rgb_mae_{role_name}"] = ( + result[rgb_abs_name] + / max(result[rgb_values_name], 1.0) + ) + + if "recall_outer" in result: + if "loss_geometry" in result: + result["loss_outer_selection"] = ( + result["loss_geometry"] + + inner_recall_weight * (1.0 - result.get("recall_inner", 1.0)) + + outer_precision_weight * (1.0 - result["precision_outer"]) + + outer_recall_weight * (1.0 - result["recall_outer"]) + + outer_iou_weight * (1.0 - result["iou_outer"]) + ) + if "hard_recall_outer" in result: + result["loss_hard_uv_selection"] = ( + inner_recall_weight * (1.0 - result["hard_iou_inner"]) + + outer_precision_weight * (1.0 - result["hard_precision_outer"]) + + outer_recall_weight * (1.0 - result["hard_recall_outer"]) + + outer_iou_weight * (1.0 - result["hard_iou_outer"]) + ) + if "hard_rgb_mae_inner" in result and "hard_rgb_mae_outer" in result: + result["loss_hard_uv_color_selection"] = ( + result["loss_hard_uv_selection"] + + hard_rgb_selection_weight + * 0.5 + * ( + result["hard_rgb_mae_inner"] + + result["hard_rgb_mae_outer"] + ) + + outer_projection_fp_selection_weight + * result.get("loss_outer_projection_false_positive", 0.0) + + outer_projection_area_selection_weight + * result.get("loss_outer_projected_area", 0.0) + ) + return result + + +def move_batch(batch, device): + return { + "uv": batch["uv"].to(device, non_blocking=True), + "path": batch["path"], + } + + +def clip_parser_gradients(model, max_norm): + """Clip the auxiliary occupancy head separately from the parser trunk.""" + parser_parameters = [] + occupancy_parameters = [] + for name, parameter in model.named_parameters(): + if parameter.grad is None: + continue + if name.startswith("outer_uv_occupancy_head."): + occupancy_parameters.append(parameter) + else: + parser_parameters.append(parameter) + if parser_parameters: + torch.nn.utils.clip_grad_norm_(parser_parameters, max_norm) + if occupancy_parameters: + torch.nn.utils.clip_grad_norm_(occupancy_parameters, max_norm) + + +def stack_view_targets(targets_by_view): + result = {} + for key in targets_by_view[0]: + stacked = torch.stack([targets[key] for targets in targets_by_view], dim=1) + if stacked.dim() == 5: + B, V, C, H, W = stacked.shape + result[key] = stacked.reshape(B * V, C, H, W) + elif stacked.dim() == 4: + B, V, H, W = stacked.shape + result[key] = stacked.reshape(B * V, H, W) + elif stacked.dim() == 3: + B, V, C = stacked.shape + result[key] = stacked.reshape(B * V, C) + else: + raise ValueError(f"Unexpected target shape for {key}: {tuple(stacked.shape)}") + return result + + +def build_parser_inputs( + batch_uv, + renderer, + views, + train, + args, + view_role_ids=None, +): + rendered_by_view = [] + targets_by_view = [] + with torch.no_grad(): + for view in views: + rendered, targets = build_dense_parser_batch( + batch_uv, + renderer, + view, + alpha_threshold=args.target_alpha_threshold, + ) + targets = dict(targets) + # Fixed-view parser training uses the canonical renderer coordinates. + # Keep an explicit identity transform for the geometry-routing API and + # for compatibility with existing affine-aware checkpoints. + targets["affine"] = rendered.new_zeros(rendered.shape[0], 3) + background_probability = args.background_augment_prob if train and args.background_augment else 0.0 + rendered = randomize_render_background( + rendered, + probability=background_probability, + bg_color=args.bg_color, + ) + rendered_by_view.append(rendered) + targets_by_view.append(targets) + + rendered = torch.stack(rendered_by_view, dim=1) + B, V, C, H, W = rendered.shape + rendered = rendered.reshape(B * V, C, H, W) + if view_role_ids is None: + view_role_ids = list(range(V)) + if len(view_role_ids) != V: + raise ValueError( + "view_role_ids must contain one role for every rendered view." + ) + view_ids = torch.tensor( + view_role_ids, + device=rendered.device, + dtype=torch.long, + ).view(1, V).expand(B, -1).reshape(B * V) + targets = stack_view_targets(targets_by_view) + return rendered, targets, V, view_ids + + +def privileged_training_views(args): + """Return primary inference views followed by training-only view groups.""" + primary = parse_views(args.views) + privileged = parse_views(getattr(args, "privileged_views", "")) + if not primary: + raise ValueError("At least one primary inference view is required.") + if privileged and len(privileged) % len(primary) != 0: + raise ValueError( + "--privileged_views must contain complete groups matching the " + f"{len(primary)} primary view roles." + ) + return primary, privileged, primary + privileged + + +def repeated_view_role_ids(primary_view_count, total_view_count): + """Map every privileged group onto the same canonical view roles.""" + if primary_view_count < 1 or total_view_count % primary_view_count != 0: + raise ValueError( + "total_view_count must contain complete primary-role groups." + ) + return [ + index % primary_view_count for index in range(total_view_count) + ] + + +def select_cached_semantic_views(features, view_indices): + """Select cached BxV semantic tensors for an active renderer view subset.""" + if features is None: + return None + + def select(value): + if value.dim() < 3: + raise ValueError( + "Cached semantic features must retain a BxV prefix." + ) + indices = torch.tensor( + view_indices, + device=value.device, + dtype=torch.long, + ) + return value.index_select(1, indices) + + if isinstance(features, dict): + return {name: select(value) for name, value in features.items()} + return select(features) + + +def visible_target_layer_uv_counts(targets, layer_index, group_size): + """Count visible GT projections of one skin layer in each grouped UV texel.""" + layer = targets["layer"] + uv = targets["uv"] + if layer.dim() != 3 or uv.dim() != 4 or uv.shape[1] != 2: + raise ValueError("Expected layer NHW and UV N2HW targets.") + if group_size <= 0: + raise ValueError("group_size must be positive.") + if layer.shape[0] % group_size != 0: + raise ValueError( + f"Target batch {layer.shape[0]} must be divisible by group_size={group_size}." + ) + + group_count = layer.shape[0] // group_size + x = (uv[:, 0].clamp(0.0, 1.0) * (UV_SIZE - 1)).round().long() + y = (uv[:, 1].clamp(0.0, 1.0) * (UV_SIZE - 1)).round().long() + flat_uv = y * UV_SIZE + x + item_group = ( + torch.arange(layer.shape[0], device=layer.device) + .div(group_size, rounding_mode="floor") + .view(-1, 1, 1) + ) + grouped_uv = flat_uv + item_group * (UV_SIZE * UV_SIZE) + visible = layer == layer_index + counts = uv.new_zeros(group_count * UV_SIZE * UV_SIZE) + selected_uv = grouped_uv[visible] + counts.scatter_add_( + 0, + selected_uv, + torch.ones(selected_uv.shape[0], device=uv.device, dtype=uv.dtype), + ) + return counts.reshape(group_count, 1, UV_SIZE, UV_SIZE) + + +def visible_target_layer_uv_mask(targets, layer_index, group_size): + return ( + visible_target_layer_uv_counts(targets, layer_index, group_size) > 0 + ).to(dtype=targets["uv"].dtype) + + +def focused_visible_layer_recall_loss( + predicted_weight, + visible_counts, + hard_fraction=0.10, + hard_weight=0.50, +): + """Combine mean recall with the worst visible texels for small-detail protection.""" + if not 0.0 < hard_fraction <= 1.0: + raise ValueError("hard_fraction must be in (0, 1].") + if not 0.0 <= hard_weight <= 1.0: + raise ValueError("hard_weight must be in [0, 1].") + if predicted_weight.shape != visible_counts.shape: + raise ValueError( + "predicted_weight and visible_counts must have identical shapes, got " + f"{tuple(predicted_weight.shape)} and {tuple(visible_counts.shape)}." + ) + + visible = visible_counts > 0 + visible_probability = ( + predicted_weight / visible_counts.clamp_min(1.0) + ).clamp(0.0, 1.0) + deficit = 1.0 - visible_probability + mean_losses = [] + hard_losses = [] + for item in range(deficit.shape[0]): + item_deficit = deficit[item][visible[item]] + if item_deficit.numel() == 0: + zero = predicted_weight[item].sum() * 0.0 + mean_losses.append(zero) + hard_losses.append(zero) + continue + mean_losses.append(item_deficit.mean()) + hard_count = max(1, math.ceil(item_deficit.numel() * hard_fraction)) + hard_losses.append(item_deficit.topk(hard_count, sorted=False).values.mean()) + + mean_loss = torch.stack(mean_losses).mean() + hard_loss = torch.stack(hard_losses).mean() + combined = (1.0 - hard_weight) * mean_loss + hard_weight * hard_loss + return combined, mean_loss, hard_loss + + +def center_weighted_texel_route_supervision( + outputs, + targets, + renderer, + views, + canonicalize=True, + center_power=2.0, +): + """Supervise the route distribution after pooling pixels into UV texels. + + The older consistency loss only makes pixels in one texel agree. It does + not decide whether the agreed label should be inner or outer. This loss + pools primary-route probabilities across all fixed views using the same + texel-center preference as inference, then applies a macro-balanced + inner/outer cross entropy to the pooled prediction. + """ + if center_power < 0.0: + raise ValueError("center_power must be non-negative.") + views = parse_views(views) + if not views: + raise ValueError("At least one view is required for texel supervision.") + if outputs["layer"].shape[1] != 3: + zero = outputs["layer"].new_zeros((), dtype=torch.float32) + return { + "loss_route_texel_supervision": zero, + "texel_route_accuracy": zero, + "texel_route_precision_outer": zero, + "texel_route_recall_outer": zero, + "texel_route_count": zero, + } + if not {"route_role", "surface", "uv"}.issubset(targets): + zero = outputs["layer"].sum() * 0.0 + return { + "loss_route_texel_supervision": zero, + "texel_route_accuracy": zero, + "texel_route_precision_outer": zero, + "texel_route_recall_outer": zero, + "texel_route_count": zero, + } + if outputs["layer"].shape[0] % len(views) != 0: + raise ValueError("Route batch must be divisible by the fixed-view count.") + + canonical_outputs = ( + canonicalize_parser_outputs(outputs) if canonicalize else outputs + ) + canonical_targets = ( + canonicalize_dense_targets(targets) if canonicalize else targets + ) + probabilities = torch.softmax( + canonical_outputs["layer"].float(), dim=1 + ) + group_count = probabilities.shape[0] // len(views) + uv_count = UV_SIZE * UV_SIZE + pooled_probability = probabilities.new_zeros( + group_count * uv_count, probabilities.shape[1] + ) + pooled_target = probabilities.new_zeros( + group_count * uv_count, probabilities.shape[1] + ) + pooled_weight = probabilities.new_zeros(group_count * uv_count, 1) + + for view_index, view in enumerate(views): + selection = slice(view_index, probabilities.shape[0], len(views)) + static = build_static_surface_routing( + renderer, view, probabilities.device + ) + route_target = canonical_targets["route_role"][selection] + surface_target = canonical_targets["surface"][selection] + uv_target = canonical_targets["uv"][selection] + foreground_target = canonical_targets["foreground"][selection, 0] + surface_count = static["texel_center_score"].shape[0] + valid_surface = ( + (surface_target >= 0) & (surface_target < surface_count) + ) + primary = ( + (route_target == 0) | (route_target == 1) + ) & valid_surface & (foreground_target > 0.5) + if not primary.any(): + continue + + safe_surface = surface_target.clamp(0, surface_count - 1) + center_score = static["texel_center_score"].unsqueeze(0).expand( + route_target.shape[0], -1, -1, -1 + ).gather(1, safe_surface.unsqueeze(1)).squeeze(1) + pixel_weight = center_score.float().clamp(0.0, 1.0).pow(center_power) + pixel_weight = pixel_weight * foreground_target.float() + primary = primary & (pixel_weight > 0.0) + if not primary.any(): + continue + + x = ( + uv_target[:, 0].clamp(0.0, 1.0) * (UV_SIZE - 1) + ).round().long() + y = ( + uv_target[:, 1].clamp(0.0, 1.0) * (UV_SIZE - 1) + ).round().long() + flat_uv = y * UV_SIZE + x + groups = torch.arange( + route_target.shape[0], device=probabilities.device + ).view(-1, 1, 1) + pooled_index = groups * uv_count + flat_uv + + selected_index = pooled_index[primary] + selected_weight = pixel_weight[primary].unsqueeze(1) + selected_probability = probabilities[selection].permute( + 0, 2, 3, 1 + )[primary] + selected_target = F.one_hot( + route_target[primary], + num_classes=probabilities.shape[1], + ).to(dtype=probabilities.dtype) + pooled_probability.index_add_( + 0, selected_index, selected_probability * selected_weight + ) + pooled_target.index_add_( + 0, selected_index, selected_target * selected_weight + ) + pooled_weight.index_add_(0, selected_index, selected_weight) + + occupied = pooled_weight[:, 0] > 0.0 + if not occupied.any(): + zero = probabilities.sum() * 0.0 + return { + "loss_route_texel_supervision": zero, + "texel_route_accuracy": zero, + "texel_route_precision_outer": zero, + "texel_route_recall_outer": zero, + "texel_route_count": zero, + } + + mean_probability = ( + pooled_probability[occupied] / pooled_weight[occupied] + ).clamp_min(1e-8) + mean_target = pooled_target[occupied] / pooled_weight[occupied] + target_role = mean_target.argmax(dim=1) + per_texel_loss = -( + mean_target * mean_probability.log() + ).sum(dim=1) + class_losses = [ + per_texel_loss[target_role == role].mean() + for role in (0, 1) + if (target_role == role).any() + ] + loss = torch.stack(class_losses).mean() + predicted_role = mean_probability.argmax(dim=1) + outer_predicted = predicted_role == 1 + outer_target = target_role == 1 + true_outer = outer_predicted & outer_target + return { + "loss_route_texel_supervision": loss, + "texel_route_accuracy": ( + predicted_role == target_role + ).float().mean(), + "texel_route_precision_outer": ( + true_outer.float().sum() + / outer_predicted.float().sum().clamp_min(1.0) + ), + "texel_route_recall_outer": ( + true_outer.float().sum() + / outer_target.float().sum().clamp_min(1.0) + ), + "texel_route_count": occupied.float().sum(), + } + + +def cross_view_outer_visibility_loss( + outputs, + target_uv, + renderer, + views, + center_power=2.0, + min_views=2, + consistency_weight=0.25, + hard_negative_fraction=0.20, + hard_negative_weight=0.75, +): + """Supervise every visible outer candidate directly from atlas alpha. + + Alpha supervision covers the union of directly visible outer texels across + all views. Only the probability-disagreement term requires a texel to be + shared by multiple views. Hard-negative mining focuses extra gradient on + transparent candidates that are most confidently misrouted as outer. + """ + views = parse_views(views) + if not views: + zero = outputs["layer"].sum() * 0.0 + return { + "loss_cross_view_outer_visibility": zero, + "loss_cross_view_outer_consistency": zero, + "loss_outer_visibility_hard_negative": zero, + "cross_view_outer_precision": zero, + "cross_view_outer_recall": zero, + "outer_visible_candidate_percent": zero, + "outer_shared_candidate_percent": zero, + "count_outer_visible_candidate_texels": zero, + "count_outer_visible_candidate_observations": zero, + "count_outer_visible_negative_candidates": zero, + "count_outer_hard_negative_candidates": zero, + "count_cross_view_outer_shared_texels": zero, + } + if min_views < 2: + raise ValueError("min_views must be at least 2.") + if not 0.0 <= consistency_weight <= 1.0: + raise ValueError("consistency_weight must be in [0, 1].") + if not 0.0 <= hard_negative_fraction <= 1.0: + raise ValueError("hard_negative_fraction must be in [0, 1].") + if hard_negative_weight < 0.0: + raise ValueError("hard_negative_weight must be non-negative.") + + probabilities = torch.softmax(outputs["layer"].float(), dim=1) + pooled, support = aggregate_direct_outer_values_by_view( + renderer, + views, + probabilities, + center_power=center_power, + ) + outer_probability = pooled[:, :, ROUTE_OUTER_PRIMARY].clamp( + 1e-6, 1.0 - 1e-6 + ) + _, outer_mask = build_uv_masks() + valid_outer = outer_mask[0].to( + device=probabilities.device, dtype=torch.bool + ).reshape(1, -1) + visible = support & valid_outer.unsqueeze(1) + visible_texels = support.any(dim=1) & valid_outer + shared = ( + support.sum(dim=1) >= int(min_views) + ) & valid_outer + target_outer = ( + target_uv[:, 3].float().reshape(target_uv.shape[0], -1) > 0.5 + ) + selected = visible + if not selected.any(): + zero = probabilities.sum() * 0.0 + return { + "loss_cross_view_outer_visibility": zero, + "loss_cross_view_outer_consistency": zero, + "loss_outer_visibility_hard_negative": zero, + "cross_view_outer_precision": zero, + "cross_view_outer_recall": zero, + "outer_visible_candidate_percent": zero, + "outer_shared_candidate_percent": zero, + "count_outer_visible_candidate_texels": zero, + "count_outer_visible_candidate_observations": zero, + "count_outer_visible_negative_candidates": zero, + "count_outer_hard_negative_candidates": zero, + "count_cross_view_outer_shared_texels": zero, + } + + expanded_target = target_outer.unsqueeze(1).expand_as(outer_probability) + per_view_loss = F.binary_cross_entropy_with_logits( + torch.logit(outer_probability[selected]), + expanded_target[selected].to(dtype=outer_probability.dtype), + reduction="none", + ) + selected_target = expanded_target[selected] + class_losses = [ + per_view_loss[selected_target == expected].mean() + for expected in (False, True) + if (selected_target == expected).any() + ] + balanced_view_loss = torch.stack(class_losses).mean() + negative_loss = per_view_loss[~selected_target] + hard_negative_count = 0 + if negative_loss.numel() > 0 and hard_negative_fraction > 0.0: + hard_negative_count = max( + 1, + math.ceil(negative_loss.numel() * hard_negative_fraction), + ) + hard_negative_loss = negative_loss.topk( + hard_negative_count, + sorted=False, + ).values.mean() + else: + hard_negative_loss = probabilities.sum() * 0.0 + + support_float = support.to(dtype=outer_probability.dtype) + consensus_probability = ( + (outer_probability * support_float).sum(dim=1) + / support_float.sum(dim=1).clamp_min(1.0) + ).clamp(1e-6, 1.0 - 1e-6) + shared_observations = support & shared.unsqueeze(1) + disagreement = ( + ( + outer_probability + - consensus_probability.unsqueeze(1) + ).square() + * support_float + * shared.unsqueeze(1) + ).sum() / shared_observations.float().sum().clamp_min(1.0) + visibility_loss = ( + balanced_view_loss + + float(hard_negative_weight) * hard_negative_loss + + float(consistency_weight) * disagreement + ) + + predicted_outer = consensus_probability[visible_texels] >= 0.5 + expected_outer = target_outer[visible_texels] + true_outer = predicted_outer & expected_outer + valid_outer_count = valid_outer.expand_as(visible_texels).float().sum() + visible_count = visible_texels.float().sum() + return { + "loss_cross_view_outer_visibility": visibility_loss, + "loss_cross_view_outer_consistency": disagreement, + "loss_outer_visibility_hard_negative": hard_negative_loss, + "cross_view_outer_precision": ( + true_outer.float().sum() + / predicted_outer.float().sum().clamp_min(1.0) + ), + "cross_view_outer_recall": ( + true_outer.float().sum() + / expected_outer.float().sum().clamp_min(1.0) + ), + "outer_visible_candidate_percent": ( + 100.0 * visible_count / valid_outer_count.clamp_min(1.0) + ), + "outer_shared_candidate_percent": ( + 100.0 + * shared.float().sum() + / visible_count.clamp_min(1.0) + ), + "count_outer_visible_candidate_texels": visible_count, + "count_outer_visible_candidate_observations": selected.float().sum(), + "count_outer_visible_negative_candidates": ( + (~selected_target).float().sum() + ), + "count_outer_hard_negative_candidates": probabilities.new_tensor( + float(hard_negative_count) + ), + "count_cross_view_outer_shared_texels": shared.float().sum(), + } + + +def privileged_view_outer_distillation_loss( + outputs, + target_uv, + renderer, + views, + primary_view_count, + center_power=2.0, + teacher_confidence=0.70, +): + """Distil training-only views into the two-view inference prediction. + + Primary and privileged views are processed as independent semantic groups + by the shared parser. Their outer-route probabilities are then projected + into the same canonical UV atlas. A privileged consensus becomes a + stop-gradient soft teacher only where it is confident and agrees with the + ground-truth outer alpha. This gives the inference views gradients from + right-side evidence without making inference depend on those images. + """ + views = parse_views(views) + if not 0.0 <= teacher_confidence <= 1.0: + raise ValueError("teacher_confidence must be in [0, 1].") + if primary_view_count < 1 or primary_view_count >= len(views): + zero = outputs["layer"].sum() * 0.0 + return { + "loss_privileged_view_distillation": zero, + "privileged_teacher_accuracy": zero, + "privileged_student_teacher_agreement": zero, + "privileged_distillation_coverage_percent": zero, + "count_privileged_distillation_texels": zero, + } + + probabilities = torch.softmax(outputs["layer"].float(), dim=1) + pooled, support = aggregate_direct_outer_values_by_view( + renderer, + views, + probabilities, + center_power=center_power, + ) + outer_probability = pooled[:, :, ROUTE_OUTER_PRIMARY].clamp( + 1e-6, 1.0 - 1e-6 + ) + primary_probability = outer_probability[:, :primary_view_count] + primary_support = support[:, :primary_view_count] + privileged_probability = outer_probability[:, primary_view_count:] + privileged_support = support[:, primary_view_count:] + + def supported_mean(values, valid): + weight = valid.to(dtype=values.dtype) + return ( + (values * weight).sum(dim=1) + / weight.sum(dim=1).clamp_min(1.0) + ) + + student_probability = supported_mean( + primary_probability, primary_support + ).clamp(1e-6, 1.0 - 1e-6) + teacher_probability = supported_mean( + privileged_probability, privileged_support + ).detach().clamp(1e-6, 1.0 - 1e-6) + student_visible = primary_support.any(dim=1) + teacher_visible = privileged_support.any(dim=1) + + _, outer_mask = build_uv_masks() + valid_outer = outer_mask[0].to( + device=probabilities.device, + dtype=torch.bool, + ).reshape(1, -1) + target_outer = ( + target_uv[:, 3].float().reshape(target_uv.shape[0], -1) > 0.5 + ) + teacher_prediction = teacher_probability >= 0.5 + teacher_correct = teacher_prediction == target_outer + confidence = torch.maximum( + teacher_probability, 1.0 - teacher_probability + ) + selected = ( + valid_outer + & student_visible + & teacher_visible + & teacher_correct + & (confidence >= float(teacher_confidence)) + ) + if not selected.any(): + zero = probabilities.sum() * 0.0 + return { + "loss_privileged_view_distillation": zero, + "privileged_teacher_accuracy": zero, + "privileged_student_teacher_agreement": zero, + "privileged_distillation_coverage_percent": zero, + "count_privileged_distillation_texels": zero, + } + + # BCE(probability, target) is intentionally rejected by PyTorch autocast. + # Convert the pooled student probability back to a stable logit and use + # the fused loss so bf16/fp16 training follows the safe kernel path. + per_texel_loss = F.binary_cross_entropy_with_logits( + torch.logit(student_probability[selected]), + teacher_probability[selected], + reduction="none", + ) + selected_target = target_outer[selected] + class_losses = [ + per_texel_loss[selected_target == expected].mean() + for expected in (False, True) + if (selected_target == expected).any() + ] + loss = torch.stack(class_losses).mean() + eligible = valid_outer & student_visible & teacher_visible + student_prediction = student_probability >= 0.5 + return { + "loss_privileged_view_distillation": loss, + "privileged_teacher_accuracy": ( + teacher_correct[eligible].float().mean() + if eligible.any() + else probabilities.sum() * 0.0 + ), + "privileged_student_teacher_agreement": ( + student_prediction[selected] == teacher_prediction[selected] + ).float().mean(), + "privileged_distillation_coverage_percent": ( + 100.0 + * selected.float().sum() + / eligible.float().sum().clamp_min(1.0) + ), + "count_privileged_distillation_texels": selected.float().sum(), + } + + +def _graph_component_labels(candidate, edge_index, max_steps=64): + """Label boolean graph components without creating gradient paths.""" + batch, node_count = candidate.shape + source, destination = edge_index + labels = torch.arange( + node_count, + device=candidate.device, + dtype=torch.long, + ).view(1, -1).expand(batch, -1).clone() + labels[~candidate] = node_count + expanded_destination = destination.view(1, -1).expand(batch, -1) + connected_edge = candidate[:, source] & candidate[:, destination] + with torch.no_grad(): + for _ in range(int(max_steps)): + neighbour_label = torch.where( + connected_edge, + labels[:, source], + torch.full_like(labels[:, source], node_count), + ) + updated = labels.clone() + updated.scatter_reduce_( + 1, + expanded_destination, + neighbour_label, + reduce="amin", + include_self=True, + ) + labels = updated + return labels + + +def outer_uv_occupancy_losses( + logits, + target_uv, + outer_part_masks, + support_mask=None, + alpha_threshold=0.5, + positive_balance=0.60, + hard_positive_fraction=0.25, + hard_negative_fraction=0.10, + topology_max_steps=64, +): + """Train projected outer occupancy with component and topology structure.""" + if logits.shape != (target_uv.shape[0], 1, UV_SIZE, UV_SIZE): + raise ValueError( + "Expected grouped outer occupancy logits shaped " + f"{(target_uv.shape[0], 1, UV_SIZE, UV_SIZE)}, got " + f"{tuple(logits.shape)}." + ) + valid = outer_part_masks[:, 0].bool().any(dim=0) + valid = valid.view(1, 1, UV_SIZE, UV_SIZE).expand_as(logits) + if support_mask is not None: + if support_mask.shape != logits.shape: + raise ValueError( + "support_mask must match occupancy logits, got " + f"{tuple(support_mask.shape)} and {tuple(logits.shape)}." + ) + # This branch owns visible route decisions. Unseen atlas texels have + # no projected evidence and are completed later, so treating them as + # supervised negatives teaches the all-transparent shortcut. + valid = valid & support_mask.to(device=logits.device, dtype=torch.bool) + target = ( + (target_uv[:, 3:4].float() > float(alpha_threshold)) + & valid + ).float() + selected_logits = logits.float()[valid] + selected_target = target[valid] + if not 0.0 <= float(positive_balance) <= 1.0: + raise ValueError("positive_balance must be in [0, 1].") + positive_mask = selected_target > 0.5 + negative_mask = ~positive_mask + positive_loss = F.softplus(-selected_logits[positive_mask]) + negative_loss = F.softplus(selected_logits[negative_mask]) + if positive_loss.numel() > 0 and negative_loss.numel() > 0: + # Average each class independently. Pixel-count BCE let transparent + # atlas texels dominate and converged to the all-empty shortcut. + loss_bce = ( + float(positive_balance) * positive_loss.mean() + + (1.0 - float(positive_balance)) * negative_loss.mean() + ) + elif positive_loss.numel() > 0: + loss_bce = positive_loss.mean() + elif negative_loss.numel() > 0: + loss_bce = negative_loss.mean() + else: + loss_bce = logits.sum() * 0.0 + probability = torch.sigmoid(logits.float()) * valid + if positive_loss.numel() > 0 and hard_positive_fraction > 0.0: + hard_positive_count = max( + 1, + math.ceil(positive_loss.numel() * hard_positive_fraction), + ) + hard_positive_loss = positive_loss.topk( + hard_positive_count, + sorted=False, + ).values.mean() + else: + hard_positive_count = 0 + hard_positive_loss = logits.sum() * 0.0 + if negative_loss.numel() > 0 and hard_negative_fraction > 0.0: + hard_negative_count = max( + 1, + math.ceil(negative_loss.numel() * hard_negative_fraction), + ) + hard_negative_loss = negative_loss.topk( + hard_negative_count, + sorted=False, + ).values.mean() + else: + hard_negative_count = 0 + hard_negative_loss = logits.sum() * 0.0 + + flat_indices, edge_index = build_outer_uv_graph() + flat_indices = flat_indices.to(logits.device) + edge_index = edge_index.to(logits.device) + probability_nodes = probability.flatten(2)[:, 0].index_select( + 1, flat_indices + ) + valid_nodes = valid.flatten(2)[:, 0].index_select(1, flat_indices) + target_nodes = target.bool().flatten(2)[:, 0].index_select( + 1, flat_indices + ) + source, destination = edge_index + valid_edges = valid_nodes[:, source] & valid_nodes[:, destination] + positive_edges = ( + valid_edges + & target_nodes[:, source] + & target_nodes[:, destination] + ) + edge_difference = ( + probability_nodes[:, source] - probability_nodes[:, destination] + ).square() + topology_consistency = ( + edge_difference[positive_edges].mean() + if positive_edges.any() + else logits.sum() * 0.0 + ) + negative_edges = ( + valid_edges + & ~target_nodes[:, source] + & ~target_nodes[:, destination] + ) + false_component_edge_score = ( + probability_nodes[:, source] * probability_nodes[:, destination] + ).square() + negative_topology = ( + false_component_edge_score[negative_edges].mean() + if negative_edges.any() + else logits.sum() * 0.0 + ) + + node_count = probability_nodes.shape[1] + component_labels = _graph_component_labels( + target_nodes, + edge_index, + max_steps=topology_max_steps, + ) + batch_offsets = ( + torch.arange(target_nodes.shape[0], device=logits.device).view(-1, 1) + * (node_count + 1) + ) + component_index = component_labels + batch_offsets + component_probability = probability_nodes.new_zeros( + target_nodes.shape[0] * (node_count + 1) + ) + component_size = probability_nodes.new_zeros( + target_nodes.shape[0] * (node_count + 1) + ) + component_probability.scatter_add_( + 0, + component_index[target_nodes], + probability_nodes[target_nodes], + ) + component_size.scatter_add_( + 0, + component_index[target_nodes], + torch.ones_like(probability_nodes[target_nodes]), + ) + existing_components = component_size > 0 + component_recall = ( + component_probability[existing_components] + / component_size[existing_components].clamp_min(1.0) + ) + component_recall_loss = ( + (1.0 - component_recall).mean() + if component_recall.numel() > 0 + else logits.sum() * 0.0 + ) + + # Macro-average false mass over predicted connected components. This gives + # a small but coherent false accessory the same component-level attention + # as a large false patch instead of letting background negatives dominate + # by raw pixel count. + predicted_nodes = (probability_nodes.detach() >= 0.5) & valid_nodes + predicted_labels = _graph_component_labels( + predicted_nodes, + edge_index, + max_steps=topology_max_steps, + ) + predicted_component_index = predicted_labels + batch_offsets + predicted_component_size = probability_nodes.new_zeros( + target_nodes.shape[0] * (node_count + 1) + ) + predicted_component_false_probability = probability_nodes.new_zeros( + target_nodes.shape[0] * (node_count + 1) + ) + predicted_component_false_count = probability_nodes.new_zeros( + target_nodes.shape[0] * (node_count + 1) + ) + predicted_component_size.scatter_add_( + 0, + predicted_component_index[predicted_nodes], + torch.ones_like(probability_nodes[predicted_nodes]), + ) + false_predicted_nodes = predicted_nodes & ~target_nodes + predicted_component_false_probability.scatter_add_( + 0, + predicted_component_index[false_predicted_nodes], + probability_nodes[false_predicted_nodes], + ) + predicted_component_false_count.scatter_add_( + 0, + predicted_component_index[false_predicted_nodes], + torch.ones_like(probability_nodes[false_predicted_nodes]), + ) + existing_predicted_components = predicted_component_size > 0 + component_false_mass = ( + predicted_component_false_probability[existing_predicted_components] + / predicted_component_size[existing_predicted_components].clamp_min(1.0) + ) + component_precision = 1.0 - ( + predicted_component_false_count[existing_predicted_components] + / predicted_component_size[existing_predicted_components].clamp_min(1.0) + ) + component_false_positive_loss = ( + component_false_mass.mean() + if component_false_mass.numel() > 0 + else logits.sum() * 0.0 + ) + dice = 1.0 - ( + 2.0 * (probability * target).sum() + 1.0 + ) / (probability.sum() + target.sum() + 1.0) + predicted = (probability >= 0.5) & valid + expected = target > 0.5 + true_positive = (predicted & expected).float().sum() + return { + "loss_outer_uv_occupancy_bce": loss_bce, + "loss_outer_uv_occupancy_dice": dice, + "loss_outer_uv_occupancy_hard_positive": hard_positive_loss, + "loss_outer_uv_occupancy_hard_negative": hard_negative_loss, + "loss_outer_topology_consistency": topology_consistency, + "loss_outer_negative_topology": negative_topology, + "loss_outer_component_recall": component_recall_loss, + "loss_outer_component_false_positive": ( + component_false_positive_loss + ), + "outer_component_recall": ( + component_recall.mean() + if component_recall.numel() > 0 + else logits.sum() * 0.0 + ), + "count_outer_components": logits.new_tensor( + float(existing_components.sum()) + ), + "count_outer_hard_positive_candidates": logits.new_tensor( + float(hard_positive_count) + ), + "count_outer_occupancy_hard_negative_candidates": logits.new_tensor( + float(hard_negative_count) + ), + "count_predicted_outer_components": logits.new_tensor( + float(existing_predicted_components.sum()) + ), + "count_outer_occupancy_supervised_texels": valid.float().sum(), + "outer_component_precision": ( + component_precision.mean() + if component_precision.numel() > 0 + else logits.sum() * 0.0 + ), + "outer_uv_occupancy_precision": ( + true_positive / predicted.float().sum().clamp_min(1.0) + ), + "outer_uv_occupancy_recall": ( + true_positive / expected.float().sum().clamp_min(1.0) + ), + } + + +def route_occupancy_agreement_loss( + outputs, + target_uv, + renderer, + views, + center_power=2.0, + confidence_threshold=0.80, +): + """Use only verified positive occupancy as an auxiliary route teacher. + + Negative occupancy is deliberately excluded: a collapsed occupancy head + otherwise teaches the route branch to erase every outer observation. + Ordinary route supervision already contains explicit outer false-positive + losses, so agreement is most useful as recall-preserving positive evidence. + """ + if "outer_uv_occupancy_logits" not in outputs: + zero = outputs["layer"].sum() * 0.0 + return { + "loss_route_occupancy_agreement": zero, + "route_occupancy_mae": zero, + "count_route_occupancy_texels": zero, + "count_route_occupancy_positive_texels": zero, + "count_route_occupancy_negative_texels": zero, + } + route_probability = torch.softmax( + outputs["layer"].float(), dim=1 + ) + pooled, support = aggregate_direct_outer_values_by_view( + renderer, + views, + route_probability, + center_power=center_power, + ) + support_float = support.to(dtype=pooled.dtype) + visible_route_outer = ( + ( + pooled[:, :, ROUTE_OUTER_PRIMARY] + * support_float + ).sum(dim=1) + / support_float.sum(dim=1).clamp_min(1.0) + ) + occupancy = torch.sigmoid( + outputs["outer_uv_occupancy_logits"].float() + ).flatten(2)[:, 0] + target_outer = ( + target_uv[:, 3].float() > 0.5 + ).flatten(1).to(device=occupancy.device) + _, outer_mask = build_uv_masks() + valid_outer = outer_mask[0].to( + device=occupancy.device, + dtype=torch.bool, + ).reshape(1, -1) + confidence_threshold = float(confidence_threshold) + confident_positive = ( + target_outer + & (occupancy.detach() >= confidence_threshold) + ) + selected = support.any(dim=1) & valid_outer & confident_positive + if not selected.any(): + zero = occupancy.sum() * 0.0 + return { + "loss_route_occupancy_agreement": zero, + "route_occupancy_mae": zero, + "count_route_occupancy_texels": zero, + "count_route_occupancy_positive_texels": zero, + "count_route_occupancy_negative_texels": zero, + } + difference = ( + visible_route_outer[selected] - occupancy.detach()[selected] + ) + return { + "loss_route_occupancy_agreement": difference.square().mean(), + "route_occupancy_mae": difference.abs().mean(), + "count_route_occupancy_texels": selected.float().sum(), + "count_route_occupancy_positive_texels": ( + selected & target_outer + ).float().sum(), + "count_route_occupancy_negative_texels": ( + selected & ~target_outer + ).float().sum(), + } + + +def differentiable_geometry_losses( + rendered, + gt_uv, + outputs, + targets, + renderer, + views, + temperature=1.0, + canonicalize=True, + recall_hard_fraction=0.10, + recall_hard_weight=0.50, + route_texel_center_power=2.0, +): + """Photometric UV and multi-view render losses for geometry-parser logits.""" + views = parse_views(views) + pred_uv, soft_details = soft_splat_geometry_predictions_to_uv( + rendered, + outputs, + renderer=renderer, + views=views, + group_size=len(views), + temperature=temperature, + canonicalize=canonicalize, + return_details=True, + ) + texel_route = center_weighted_texel_route_supervision( + outputs, + targets, + renderer, + views, + canonicalize=canonicalize, + center_power=route_texel_center_power, + ) + gt_uv = gt_uv.float() + support = (pred_uv[:, 3:4].detach() > 0.05).to(dtype=pred_uv.dtype) + gt_alpha = (gt_uv[:, 3:4] > 0.5).to(dtype=pred_uv.dtype) + rgb_mask = support * gt_alpha + rgb_denom = (rgb_mask.sum(dim=(1, 2, 3)) * 3.0).clamp_min(1.0) + loss_soft_uv_rgb = ( + ((pred_uv[:, :3] - gt_uv[:, :3]).abs() * rgb_mask) + .sum(dim=(1, 2, 3)) + .div(rgb_denom) + .mean() + ) + transparent_support = support * (1.0 - gt_alpha) + alpha_denom = transparent_support.sum(dim=(1, 2, 3)).clamp_min(1.0) + loss_soft_uv_alpha = ( + (pred_uv[:, 3:4] * transparent_support) + .sum(dim=(1, 2, 3)) + .div(alpha_denom) + .mean() + ) + visible_inner_counts = visible_target_layer_uv_counts( + targets, + layer_index=0, + group_size=len(views), + ).to(dtype=pred_uv.dtype) + visible_inner_counts = visible_inner_counts * gt_alpha + visible_inner_uv = (visible_inner_counts > 0).to(dtype=pred_uv.dtype) + ( + loss_soft_uv_inner_recall, + loss_soft_uv_inner_recall_mean, + loss_soft_uv_inner_recall_hard, + ) = focused_visible_layer_recall_loss( + soft_details["layer_weight"][:, 0:1], + visible_inner_counts, + hard_fraction=recall_hard_fraction, + hard_weight=recall_hard_weight, + ) + + visible_outer_counts = visible_target_layer_uv_counts( + targets, + layer_index=1, + group_size=len(views), + ).to(dtype=pred_uv.dtype) + visible_outer_counts = visible_outer_counts * gt_alpha + visible_outer_uv = (visible_outer_counts > 0).to(dtype=pred_uv.dtype) + ( + loss_soft_uv_outer_recall, + loss_soft_uv_outer_recall_mean, + loss_soft_uv_outer_recall_hard, + ) = focused_visible_layer_recall_loss( + soft_details["layer_weight"][:, 1:2], + visible_outer_counts, + hard_fraction=recall_hard_fraction, + hard_weight=recall_hard_weight, + ) + + render_rgb_total = pred_uv.new_zeros(()) + render_alpha_total = pred_uv.new_zeros(()) + outer_projection_false_positive_total = pred_uv.new_zeros(()) + outer_projection_false_negative_total = pred_uv.new_zeros(()) + outer_projection_dice_total = pred_uv.new_zeros(()) + outer_projected_area_total = pred_uv.new_zeros(()) + outer_projected_false_positive_pixels = pred_uv.new_zeros(()) + outer_projected_negative_pixels = pred_uv.new_zeros(()) + for view in views: + pred_render = render_direct_uv(pred_uv, renderer, view) + with torch.no_grad(): + gt_render = render_direct_uv(gt_uv, renderer, view) + gt_outer_alpha = render_outer_alpha_uv(gt_uv, renderer, view) + pred_outer_alpha = render_outer_alpha_uv( + pred_uv, renderer, view + ).clamp(1e-6, 1.0 - 1e-6) + outer_projection_mask = getattr( + renderer, f"{view}_outer_mask" + ).to( + device=pred_uv.device, + dtype=pred_uv.dtype, + ).view(1, 1, *pred_outer_alpha.shape[-2:]) + outer_negative = ( + outer_projection_mask - gt_outer_alpha.detach() + ).clamp(0.0, 1.0) + outer_positive = gt_outer_alpha.detach().clamp(0.0, 1.0) + outer_negative_denom = outer_negative.sum( + dim=(1, 2, 3) + ).clamp_min(1.0) + outer_positive_denom = outer_positive.sum( + dim=(1, 2, 3) + ).clamp_min(1.0) + outer_projection_false_positive_total = ( + outer_projection_false_positive_total + + ( + -torch.log1p(-pred_outer_alpha) + * outer_negative + ).sum(dim=(1, 2, 3)).div(outer_negative_denom).mean() + ) + outer_projection_false_negative_total = ( + outer_projection_false_negative_total + + ( + -torch.log(pred_outer_alpha) * outer_positive + ).sum(dim=(1, 2, 3)).div(outer_positive_denom).mean() + ) + outer_projection_dice_total = ( + outer_projection_dice_total + + ( + 1.0 + - ( + 2.0 + * (pred_outer_alpha * outer_positive).sum( + dim=(1, 2, 3) + ) + + 1.0 + ) + / ( + (pred_outer_alpha * outer_projection_mask).sum( + dim=(1, 2, 3) + ) + + outer_positive.sum(dim=(1, 2, 3)) + + 1.0 + ) + ).mean() + ) + predicted_outer_area = ( + pred_outer_alpha * outer_projection_mask + ).sum(dim=(1, 2, 3)) + expected_outer_area = outer_positive.sum(dim=(1, 2, 3)) + outer_projected_area_total = ( + outer_projected_area_total + + ( + predicted_outer_area - expected_outer_area + ).abs().div(expected_outer_area.clamp_min(1.0)).mean() + ) + outer_projected_false_positive_pixels = ( + outer_projected_false_positive_pixels + + ( + (pred_outer_alpha.detach() >= 0.5) + & (outer_positive < 0.5) + & (outer_projection_mask > 0.5) + ).float().sum() + ) + outer_projected_negative_pixels = ( + outer_projected_negative_pixels + + ( + (outer_positive < 0.5) + & (outer_projection_mask > 0.5) + ).float().sum() + ) + foreground = gt_render[:, 3:4].detach() + rgb_denom = (foreground.sum(dim=(1, 2, 3)) * 3.0).clamp_min(1.0) + render_rgb_total = render_rgb_total + ( + ((pred_render[:, :3] - gt_render[:, :3]).abs() * foreground) + .sum(dim=(1, 2, 3)) + .div(rgb_denom) + .mean() + ) + alpha_support = torch.maximum( + pred_render[:, 3:4], gt_render[:, 3:4] + ).detach() + alpha_denom = alpha_support.sum(dim=(1, 2, 3)).clamp_min(1.0) + render_alpha_total = render_alpha_total + ( + ((pred_render[:, 3:4] - gt_render[:, 3:4]).abs() * alpha_support) + .sum(dim=(1, 2, 3)) + .div(alpha_denom) + .mean() + ) + + view_count = max(len(views), 1) + return { + **texel_route, + "loss_soft_uv_rgb": loss_soft_uv_rgb, + "loss_soft_uv_alpha": loss_soft_uv_alpha, + "loss_soft_uv_inner_recall": loss_soft_uv_inner_recall, + "loss_soft_uv_inner_recall_mean": loss_soft_uv_inner_recall_mean, + "loss_soft_uv_inner_recall_hard": loss_soft_uv_inner_recall_hard, + "loss_soft_uv_outer_recall": loss_soft_uv_outer_recall, + "loss_soft_uv_outer_recall_mean": loss_soft_uv_outer_recall_mean, + "loss_soft_uv_outer_recall_hard": loss_soft_uv_outer_recall_hard, + "loss_render_rgb": render_rgb_total / view_count, + "loss_render_alpha": render_alpha_total / view_count, + "loss_outer_projection_false_positive": ( + outer_projection_false_positive_total / view_count + ), + "loss_outer_projection_false_negative": ( + outer_projection_false_negative_total / view_count + ), + "loss_outer_projection_dice": ( + outer_projection_dice_total / view_count + ), + "loss_outer_projected_area": ( + outer_projected_area_total / view_count + ), + "outer_projected_false_positive_rate": ( + outer_projected_false_positive_pixels + / outer_projected_negative_pixels.clamp_min(1.0) + ), + "count_outer_projected_false_positive_pixels": ( + outer_projected_false_positive_pixels + ), + "count_outer_projected_negative_pixels": ( + outer_projected_negative_pixels + ), + "visible_inner_uv_percent": visible_inner_uv.float().mean() * 100.0, + "visible_outer_uv_percent": visible_outer_uv.float().mean() * 100.0, + "soft_uv_known_percent": ( + (soft_details["support"] > 0.05).float().mean() * 100.0 + ), + } + + +def hard_uv_conditioning_metrics( + rendered, + outputs, + targets, + renderer, + views, + args, +): + """Measure known-texel precision/recall through the exact hard inference route.""" + group_size = len(views) + hard_outputs = { + name: ( + value.to(dtype=rendered.dtype) + if torch.is_tensor(value) and value.is_floating_point() + else value + ) + for name, value in outputs.items() + } + if "affine" in outputs: + predicted = splat_parser_predictions_to_uv_conditioning( + rendered, + hard_outputs, + renderer=renderer, + views=views, + group_size=group_size, + fg_threshold=args.splat_fg_threshold, + bg_color=args.bg_color, + semantic_gate=args.semantic_gate, + affine_refine=args.affine_refine, + affine_refine_translation_px=args.affine_refine_translation_px, + affine_refine_scale=args.affine_refine_scale, + route_confidence_threshold=args.route_confidence_threshold, + route_margin_threshold=args.route_margin_threshold, + outer_route_confidence_threshold=args.outer_route_confidence_threshold, + outer_route_margin_threshold=args.outer_route_margin_threshold, + outer_uv_min_coverage=args.outer_uv_min_coverage, + outer_uv_min_source_pixels=args.outer_uv_min_source_pixels, + outer_silhouette_consistency=args.outer_silhouette_consistency, + outer_silhouette_min_coverage=( + args.outer_silhouette_min_coverage + ), + outer_silhouette_dilation=args.outer_silhouette_dilation, + outer_silhouette_min_pixels=args.outer_silhouette_min_pixels, + outer_geometry_rescue=getattr(args, "outer_geometry_rescue", False), + outer_rescue_confidence_threshold=getattr( + args, "outer_rescue_confidence_threshold", 0.60 + ), + outer_rescue_margin_threshold=getattr( + args, "outer_rescue_margin_threshold", 0.25 + ), + outer_rescue_min_coverage=getattr( + args, "outer_rescue_min_coverage", 0.10 + ), + color_aggregation=args.splat_color_aggregation, + geometry_route_texel_consensus=getattr( + args, "geometry_route_texel_consensus", False + ), + geometry_route_texel_consensus_weight=getattr( + args, "geometry_route_texel_consensus_weight", 0.60 + ), + geometry_route_preserve_outer_confidence=getattr( + args, "geometry_route_preserve_outer_confidence", 0.80 + ), + geometry_route_preserve_outer_margin=getattr( + args, "geometry_route_preserve_outer_margin", 0.35 + ), + geometry_route_consensus_outer_confidence=getattr( + args, "geometry_route_consensus_outer_confidence", 0.70 + ), + geometry_route_consensus_outer_margin=getattr( + args, "geometry_route_consensus_outer_margin", 0.20 + ), + geometry_cross_view_outer_consistency=getattr( + args, "geometry_cross_view_outer_consistency", False + ), + geometry_cross_view_outer_weight=getattr( + args, "geometry_cross_view_outer_weight", 0.50 + ), + geometry_cross_view_outer_positive_confidence=getattr( + args, + "geometry_cross_view_outer_positive_confidence", + 0.70, + ), + geometry_cross_view_outer_positive_margin=getattr( + args, "geometry_cross_view_outer_positive_margin", 0.20 + ), + geometry_cross_view_outer_negative_confidence=getattr( + args, + "geometry_cross_view_outer_negative_confidence", + 0.70, + ), + geometry_cross_view_outer_negative_margin=getattr( + args, "geometry_cross_view_outer_negative_margin", 0.20 + ), + geometry_cross_view_outer_background_max_coverage=getattr( + args, + "geometry_cross_view_outer_background_max_coverage", + 0.25, + ), + geometry_cross_view_outer_min_views=getattr( + args, "geometry_cross_view_outer_min_views", 2 + ), + outer_uv_occupancy=getattr( + args, "outer_uv_occupancy_routing", False + ), + outer_uv_occupancy_blend_weight=getattr( + args, "outer_uv_occupancy_blend_weight", 0.30 + ), + outer_uv_occupancy_gate_threshold=getattr( + args, "outer_uv_occupancy_gate_threshold", 0.15 + ), + outer_uv_occupancy_rescue_threshold=getattr( + args, "outer_uv_occupancy_rescue_threshold", 0.70 + ), + outer_uv_occupancy_rescue_route_threshold=getattr( + args, + "outer_uv_occupancy_rescue_route_threshold", + 0.30, + ), + outer_uv_component_routing=getattr( + args, "outer_uv_component_routing", False + ), + outer_uv_component_seed_threshold=getattr( + args, "outer_uv_component_seed_threshold", 0.80 + ), + outer_uv_component_grow_threshold=getattr( + args, "outer_uv_component_grow_threshold", 0.50 + ), + outer_uv_component_min_size=getattr( + args, "outer_uv_component_min_size", 2 + ), + observed_foreground=None, + background_color_tolerance=args.background_color_tolerance, + color_background_tolerance=getattr( + args, "color_background_tolerance", 8.0 / 255.0 + ), + color_foreground_inset=getattr(args, "color_foreground_inset", 1), + reject_semantic_fallback=not args.allow_semantic_fallback, + ) + expected = splat_deterministic_targets_to_uv_conditioning( + rendered, + targets, + renderer=renderer, + views=views, + group_size=group_size, + bg_color=args.bg_color, + color_aggregation=args.splat_color_aggregation, + ) + else: + predicted = splat_predictions_to_uv_conditioning( + rendered, + hard_outputs, + group_size=group_size, + fg_threshold=args.splat_fg_threshold, + bg_color=args.bg_color, + ) + expected = splat_targets_to_uv_conditioning( + rendered, + targets, + group_size=group_size, + bg_color=args.bg_color, + ) + + metrics = {} + for role_name, rgba_start, known_channel in ( + ("inner", 0, 4), + ("outer", 5, 9), + ): + predicted_known = predicted[:, known_channel] > 0.5 + expected_known = expected[:, known_channel] > 0.5 + matched = predicted_known & expected_known + metrics[f"count_hard_{role_name}_tp"] = matched.sum().float() + metrics[f"count_hard_{role_name}_fp"] = ( + predicted_known & ~expected_known + ).sum().float() + metrics[f"count_hard_{role_name}_fn"] = ( + ~predicted_known & expected_known + ).sum().float() + rgb_error = ( + predicted[:, rgba_start : rgba_start + 3] + - expected[:, rgba_start : rgba_start + 3] + ).abs() + metrics[f"count_hard_{role_name}_rgb_abs"] = ( + rgb_error * matched.unsqueeze(1) + ).sum().float() + metrics[f"count_hard_{role_name}_rgb_values"] = ( + matched.sum().float() * 3.0 + ) + return metrics + + +def run_epoch( + model, + criterion, + renderer, + loader, + optimizer, + scaler, + device, + precision, + args, + train=True, + compute_hard_metrics=False, + semantic_cache=None, + semantic_masks=None, +): + model.train(train) + primary_views, privileged_views, all_training_views = ( + privileged_training_views(args) + ) + # Validation and checkpoint selection must exactly match inference. Only + # the training pass receives the privileged right-side view group. + views = all_training_views if train else primary_views + view_role_ids = repeated_view_role_ids( + len(primary_views), len(views) + ) + metric_sums = {} + sample_count = 0 + iterator = tqdm(loader, leave=False, file=sys.__stderr__ or sys.stderr) if tqdm is not None else loader + if train: + optimizer.zero_grad(set_to_none=True) + + for batch_index, batch in enumerate(iterator): + batch = move_batch(batch, device) + rendered, targets, _, view_ids = build_parser_inputs( + batch["uv"], + renderer, + views, + train=train, + args=args, + view_role_ids=view_role_ids, + ) + parser_samples = rendered.shape[0] + semantic_features = cached_semantic_batch( + semantic_cache, batch["path"], device + ) + if semantic_features is not None and not train and privileged_views: + semantic_features = select_cached_semantic_views( + semantic_features, + range(len(primary_views)), + ) + + with torch.set_grad_enabled(train): + with autocast_context(device, precision): + semantic_kwargs = {} + if semantic_features is not None: + semantic_kwargs["semantic_features"] = semantic_features + elif getattr(args, "semantic_backbone", "none") != "none": + semantic_kwargs["semantic_foreground"] = targets["foreground"] + outputs = model( + rendered, + view_ids=view_ids, + **semantic_kwargs, + ) + outputs = attach_projected_outer_uv_occupancy( + model, + outputs, + renderer, + views, + observed_foreground=targets["foreground"][:, 0], + center_power=args.route_texel_center_power, + ) + losses = criterion(outputs, targets) + zero = losses["loss_total"].new_zeros(()) + if semantic_masks is not None and "outer_presence_logits" in outputs: + inner_part_masks, outer_part_masks = semantic_masks + attributes = build_semantic_attribute_targets( + batch["uv"], inner_part_masks, outer_part_masks + ) + semantic_group_repeats = ( + outputs["outer_presence_logits"].shape[0] + // attributes["outer_presence"].shape[0] + ) + if semantic_group_repeats > 1: + attributes = { + name: value.repeat_interleave( + semantic_group_repeats, dim=0 + ) + for name, value in attributes.items() + } + loss_semantic_presence = F.binary_cross_entropy_with_logits( + outputs["outer_presence_logits"].float(), + attributes["outer_presence"], + ) + loss_semantic_coverage = F.smooth_l1_loss( + outputs["outer_coverage"].float(), + attributes["outer_coverage"], + ) + loss_semantic_attributes = ( + args.lambda_semantic_presence * loss_semantic_presence + + args.lambda_semantic_coverage * loss_semantic_coverage + ) + losses["loss_semantic_presence"] = loss_semantic_presence + losses["loss_semantic_coverage"] = loss_semantic_coverage + losses["loss_semantic_attributes"] = loss_semantic_attributes + losses["loss_total"] = losses["loss_total"] + loss_semantic_attributes + losses["loss_routing"] = losses["loss_routing"] + loss_semantic_attributes + else: + losses["loss_semantic_presence"] = zero + losses["loss_semantic_coverage"] = zero + losses["loss_semantic_attributes"] = zero + if ( + semantic_masks is not None + and "outer_uv_occupancy_logits" in outputs + ): + _, outer_part_masks = semantic_masks + occupancy = outer_uv_occupancy_losses( + outputs["outer_uv_occupancy_logits"], + batch["uv"], + outer_part_masks, + support_mask=outputs.get( + "outer_uv_occupancy_support" + ), + alpha_threshold=args.target_alpha_threshold, + positive_balance=( + args.outer_uv_occupancy_positive_balance + ), + hard_positive_fraction=( + args.outer_hard_positive_fraction + ), + hard_negative_fraction=( + args.outer_hard_negative_fraction + ), + ) + loss_outer_uv_occupancy = ( + occupancy["loss_outer_uv_occupancy_bce"] + + args.outer_uv_occupancy_dice_weight + * occupancy["loss_outer_uv_occupancy_dice"] + + args.outer_hard_positive_weight + * occupancy[ + "loss_outer_uv_occupancy_hard_positive" + ] + + args.outer_hard_negative_weight + * occupancy[ + "loss_outer_uv_occupancy_hard_negative" + ] + ) + weighted_outer_uv_occupancy = ( + args.lambda_outer_uv_occupancy + * loss_outer_uv_occupancy + + args.lambda_outer_component_recall + * occupancy["loss_outer_component_recall"] + + args.lambda_outer_component_false_positive + * occupancy[ + "loss_outer_component_false_positive" + ] + + args.lambda_outer_topology + * occupancy["loss_outer_topology_consistency"] + + args.lambda_outer_negative_topology + * occupancy["loss_outer_negative_topology"] + ) + losses.update(occupancy) + losses["loss_outer_uv_occupancy"] = ( + loss_outer_uv_occupancy + ) + losses["loss_outer_uv_occupancy_weighted"] = ( + weighted_outer_uv_occupancy + ) + losses["loss_total"] = ( + losses["loss_total"] + weighted_outer_uv_occupancy + ) + losses["loss_routing"] = ( + losses["loss_routing"] + weighted_outer_uv_occupancy + ) + else: + losses["loss_outer_uv_occupancy_bce"] = zero + losses["loss_outer_uv_occupancy_dice"] = zero + losses["loss_outer_uv_occupancy"] = zero + losses["loss_outer_uv_occupancy_weighted"] = zero + losses["outer_uv_occupancy_precision"] = zero + losses["outer_uv_occupancy_recall"] = zero + losses["loss_outer_uv_occupancy_hard_positive"] = zero + losses["loss_outer_uv_occupancy_hard_negative"] = zero + losses["loss_outer_topology_consistency"] = zero + losses["loss_outer_negative_topology"] = zero + losses["loss_outer_component_recall"] = zero + losses["loss_outer_component_false_positive"] = zero + losses["outer_component_recall"] = zero + losses["outer_component_precision"] = zero + losses["count_outer_components"] = zero + losses["count_predicted_outer_components"] = zero + losses["count_outer_occupancy_supervised_texels"] = zero + losses["count_outer_hard_positive_candidates"] = zero + losses[ + "count_outer_occupancy_hard_negative_candidates" + ] = zero + + agreement = route_occupancy_agreement_loss( + outputs, + batch["uv"], + renderer, + views, + center_power=args.route_texel_center_power, + confidence_threshold=( + args.outer_occupancy_agreement_confidence_threshold + ), + ) + if train: + batch_progress = float(batch_index) / max( + len(loader) - 1, + 1, + ) + warmup = float( + args.outer_occupancy_agreement_warmup_fraction + ) + agreement_scale = max( + 0.0, + min( + 1.0, + (batch_progress - warmup) + / max(warmup, 1e-6), + ), + ) + else: + agreement_scale = 1.0 + weighted_agreement = ( + args.lambda_route_occupancy_agreement + * agreement_scale + * agreement["loss_route_occupancy_agreement"] + ) + losses.update(agreement) + losses["route_occupancy_agreement_scale"] = zero.new_tensor( + agreement_scale + ) + losses["loss_route_occupancy_agreement_weighted"] = ( + weighted_agreement + ) + losses["loss_total"] = losses["loss_total"] + weighted_agreement + losses["loss_routing"] = ( + losses["loss_routing"] + weighted_agreement + ) + auxiliary_enabled = args.parser_mode == "geometry_fit" and any( + weight > 0 + for weight in ( + args.lambda_soft_uv_rgb, + args.lambda_soft_uv_alpha, + args.lambda_soft_uv_inner_recall, + args.lambda_soft_uv_outer_recall, + args.lambda_render_rgb, + args.lambda_render_alpha, + args.lambda_outer_projection_false_positive, + args.lambda_outer_projection_false_negative, + args.lambda_outer_projection_dice, + args.lambda_outer_projected_area, + args.lambda_route_texel_supervision, + args.lambda_cross_view_outer_visibility, + ) + ) + if auxiliary_enabled: + auxiliary = differentiable_geometry_losses( + rendered, + batch["uv"], + outputs, + targets, + renderer, + views, + temperature=args.render_softmax_temperature, + canonicalize=False, + recall_hard_fraction=args.soft_uv_recall_hard_fraction, + recall_hard_weight=args.soft_uv_recall_hard_weight, + route_texel_center_power=( + args.route_texel_center_power + ), + ) + cross_view_visibility = cross_view_outer_visibility_loss( + outputs, + batch["uv"], + renderer, + views, + center_power=args.route_texel_center_power, + min_views=args.geometry_cross_view_outer_min_views, + consistency_weight=( + args.cross_view_outer_consistency_loss_weight + ), + hard_negative_fraction=( + args.outer_visibility_hard_negative_fraction + ), + hard_negative_weight=( + args.outer_visibility_hard_negative_weight + ), + ) + auxiliary.update(cross_view_visibility) + else: + auxiliary = { + "loss_soft_uv_rgb": zero, + "loss_soft_uv_alpha": zero, + "loss_soft_uv_inner_recall": zero, + "loss_soft_uv_inner_recall_mean": zero, + "loss_soft_uv_inner_recall_hard": zero, + "loss_soft_uv_outer_recall": zero, + "loss_soft_uv_outer_recall_mean": zero, + "loss_soft_uv_outer_recall_hard": zero, + "loss_render_rgb": zero, + "loss_render_alpha": zero, + "loss_outer_projection_false_positive": zero, + "loss_outer_projection_false_negative": zero, + "loss_outer_projection_dice": zero, + "loss_outer_projected_area": zero, + "outer_projected_false_positive_rate": zero, + "count_outer_projected_false_positive_pixels": zero, + "count_outer_projected_negative_pixels": zero, + "loss_route_texel_supervision": zero, + "texel_route_accuracy": zero, + "texel_route_precision_outer": zero, + "texel_route_recall_outer": zero, + "texel_route_count": zero, + "loss_cross_view_outer_visibility": zero, + "loss_cross_view_outer_consistency": zero, + "loss_outer_visibility_hard_negative": zero, + "cross_view_outer_precision": zero, + "cross_view_outer_recall": zero, + "outer_visible_candidate_percent": zero, + "outer_shared_candidate_percent": zero, + "count_outer_visible_candidate_texels": zero, + "count_outer_visible_candidate_observations": zero, + "count_outer_visible_negative_candidates": zero, + "count_outer_hard_negative_candidates": zero, + "count_cross_view_outer_shared_texels": zero, + "soft_uv_known_percent": zero, + "visible_inner_uv_percent": zero, + "visible_outer_uv_percent": zero, + } + if train and privileged_views: + privileged_distillation = ( + privileged_view_outer_distillation_loss( + outputs, + batch["uv"], + renderer, + views, + primary_view_count=len(primary_views), + center_power=args.route_texel_center_power, + teacher_confidence=( + args.privileged_view_teacher_confidence + ), + ) + ) + else: + privileged_distillation = { + "loss_privileged_view_distillation": zero, + "privileged_teacher_accuracy": zero, + "privileged_student_teacher_agreement": zero, + "privileged_distillation_coverage_percent": zero, + "count_privileged_distillation_texels": zero, + } + auxiliary.update(privileged_distillation) + weighted_auxiliary = ( + args.lambda_soft_uv_rgb * auxiliary["loss_soft_uv_rgb"] + + args.lambda_soft_uv_alpha * auxiliary["loss_soft_uv_alpha"] + + args.lambda_soft_uv_inner_recall + * auxiliary["loss_soft_uv_inner_recall"] + + args.lambda_soft_uv_outer_recall + * auxiliary["loss_soft_uv_outer_recall"] + + args.lambda_render_rgb * auxiliary["loss_render_rgb"] + + args.lambda_render_alpha * auxiliary["loss_render_alpha"] + + args.lambda_outer_projection_false_positive + * auxiliary["loss_outer_projection_false_positive"] + + args.lambda_outer_projection_false_negative + * auxiliary["loss_outer_projection_false_negative"] + + args.lambda_outer_projection_dice + * auxiliary["loss_outer_projection_dice"] + + args.lambda_outer_projected_area + * auxiliary["loss_outer_projected_area"] + + args.lambda_route_texel_supervision + * auxiliary["loss_route_texel_supervision"] + + args.lambda_cross_view_outer_visibility + * auxiliary["loss_cross_view_outer_visibility"] + + args.lambda_privileged_view_distillation + * auxiliary["loss_privileged_view_distillation"] + ) + losses.update(auxiliary) + losses["loss_differentiable"] = weighted_auxiliary + losses["loss_total"] = losses["loss_total"] + weighted_auxiliary + losses["loss_geometry"] = losses["loss_geometry"] + weighted_auxiliary + losses["loss_routing"] = losses["loss_routing"] + weighted_auxiliary + loss = losses["loss_total"] + + if train: + final_group_size = len(loader) % args.gradient_accumulation_steps + accumulation_denominator = args.gradient_accumulation_steps + if ( + final_group_size + and batch_index >= len(loader) - final_group_size + ): + accumulation_denominator = final_group_size + accumulated_loss = loss / float( + accumulation_denominator + ) + should_step = ( + (batch_index + 1) % args.gradient_accumulation_steps == 0 + or batch_index + 1 == len(loader) + ) + if scaler is not None: + scaler.scale(accumulated_loss).backward() + if should_step: + scaler.unscale_(optimizer) + clip_parser_gradients(model, args.grad_clip) + scaler.step(optimizer) + scaler.update() + else: + accumulated_loss.backward() + if should_step: + clip_parser_gradients(model, args.grad_clip) + optimizer.step() + if should_step: + optimizer.zero_grad(set_to_none=True) + + if compute_hard_metrics: + with torch.no_grad(): + losses.update( + hard_uv_conditioning_metrics( + rendered, + outputs, + targets, + renderer, + views, + args, + ) + ) + + sample_count += parser_samples + for name, value in losses.items(): + detached = value.detach() + metric_weight = 1 if name.startswith("count_") else parser_samples + metric_sums[name] = ( + metric_sums.get(name, detached.new_zeros(())) + + detached * metric_weight + ) + + if tqdm is not None and args.log_every > 0 and sample_count % (args.log_every * parser_samples) == 0: + avg = format_metrics( + metric_sums, + sample_count, + inner_recall_weight=args.inner_selection_recall_weight, + outer_precision_weight=args.outer_selection_precision_weight, + outer_recall_weight=args.outer_selection_recall_weight, + outer_iou_weight=args.outer_selection_iou_weight, + hard_rgb_selection_weight=args.hard_rgb_selection_weight, + outer_projection_fp_selection_weight=( + args.outer_projection_fp_selection_weight + ), + outer_projection_area_selection_weight=( + args.outer_projection_area_selection_weight + ), + ) + postfix = { + "total": f"{avg['loss_total']:.4f}", + "fg": f"{avg.get('recall_foreground', avg['acc_foreground']):.3f}", + } + if args.parser_mode in ("global_affine", "geometry_fit"): + postfix["align"] = f"{avg.get('err_affine_translation_px', 0.0):.2f}px" + postfix["scale"] = f"{avg.get('err_affine_scale_pct', 0.0):.2f}%" + if args.parser_mode == "geometry_fit": + postfix["outer_p"] = f"{avg.get('precision_outer', 0.0):.3f}" + postfix["outer_r"] = f"{avg.get('recall_outer', 0.0):.3f}" + postfix["render"] = f"{avg.get('loss_render_rgb', 0.0):.3f}" + else: + postfix["surface"] = f"{avg.get('acc_surface', 0.0):.3f}" + else: + postfix["uv"] = f"{avg.get('loss_uv_l1_px', avg['loss_uv']):.2f}px" + postfix["uv1"] = f"{avg.get('acc_uv_within1', 0.0):.3f}" + iterator.set_postfix( + **postfix, + ) + + return format_metrics( + metric_sums, + sample_count, + inner_recall_weight=args.inner_selection_recall_weight, + outer_precision_weight=args.outer_selection_precision_weight, + outer_recall_weight=args.outer_selection_recall_weight, + outer_iou_weight=args.outer_selection_iou_weight, + hard_rgb_selection_weight=args.hard_rgb_selection_weight, + outer_projection_fp_selection_weight=( + args.outer_projection_fp_selection_weight + ), + outer_projection_area_selection_weight=( + args.outer_projection_area_selection_weight + ), + ) + + + + + +def save_preview( + model, + renderer, + loader, + device, + args, + output_path, + max_items=2, + semantic_cache=None, +): + model.eval() + views = parse_views(args.views) + batch = move_batch(next(iter(loader)), device) + rendered, targets, view_count, view_ids = build_parser_inputs( + batch["uv"], renderer, views, train=False, args=args + ) + semantic_features = cached_semantic_batch( + semantic_cache, batch["path"], device + ) + if semantic_features is not None and parse_views( + getattr(args, "privileged_views", "") + ): + semantic_features = select_cached_semantic_views( + semantic_features, + range(len(views)), + ) + with torch.no_grad(): + semantic_kwargs = {} + if semantic_features is not None: + semantic_kwargs["semantic_features"] = semantic_features + elif getattr(args, "semantic_backbone", "none") != "none": + semantic_kwargs["semantic_foreground"] = targets["foreground"] + outputs = model( + rendered, + view_ids=view_ids, + **semantic_kwargs, + ) + outputs = attach_projected_outer_uv_occupancy( + model, + outputs, + renderer, + views, + observed_foreground=targets["foreground"][:, 0], + center_power=getattr(args, "route_texel_center_power", 2.0), + ) + if "affine" in outputs: + pred_conditioning, routing_details = splat_parser_predictions_to_uv_conditioning( + rendered, + outputs, + renderer=renderer, + views=views, + group_size=view_count, + fg_threshold=args.splat_fg_threshold, + bg_color=args.bg_color, + semantic_gate=args.semantic_gate, + affine_refine=args.affine_refine, + affine_refine_translation_px=args.affine_refine_translation_px, + affine_refine_scale=args.affine_refine_scale, + route_confidence_threshold=args.route_confidence_threshold, + route_margin_threshold=args.route_margin_threshold, + outer_route_confidence_threshold=args.outer_route_confidence_threshold, + outer_route_margin_threshold=args.outer_route_margin_threshold, + outer_uv_min_coverage=args.outer_uv_min_coverage, + outer_uv_min_source_pixels=getattr( + args, "outer_uv_min_source_pixels", 15 + ), + outer_silhouette_consistency=getattr( + args, "outer_silhouette_consistency", True + ), + outer_silhouette_min_coverage=getattr( + args, "outer_silhouette_min_coverage", 0.50 + ), + outer_silhouette_dilation=getattr( + args, "outer_silhouette_dilation", 0 + ), + outer_silhouette_min_pixels=getattr( + args, "outer_silhouette_min_pixels", 4 + ), + outer_geometry_rescue=getattr(args, "outer_geometry_rescue", False), + outer_rescue_confidence_threshold=getattr( + args, "outer_rescue_confidence_threshold", 0.60 + ), + outer_rescue_margin_threshold=getattr( + args, "outer_rescue_margin_threshold", 0.25 + ), + outer_rescue_min_coverage=getattr( + args, "outer_rescue_min_coverage", 0.10 + ), + color_aggregation=args.splat_color_aggregation, + geometry_route_texel_consensus=getattr( + args, "geometry_route_texel_consensus", False + ), + geometry_route_texel_consensus_weight=getattr( + args, "geometry_route_texel_consensus_weight", 0.60 + ), + geometry_route_preserve_outer_confidence=getattr( + args, "geometry_route_preserve_outer_confidence", 0.80 + ), + geometry_route_preserve_outer_margin=getattr( + args, "geometry_route_preserve_outer_margin", 0.35 + ), + geometry_route_consensus_outer_confidence=getattr( + args, "geometry_route_consensus_outer_confidence", 0.70 + ), + geometry_route_consensus_outer_margin=getattr( + args, "geometry_route_consensus_outer_margin", 0.20 + ), + geometry_cross_view_outer_consistency=getattr( + args, "geometry_cross_view_outer_consistency", False + ), + geometry_cross_view_outer_weight=getattr( + args, "geometry_cross_view_outer_weight", 0.50 + ), + geometry_cross_view_outer_positive_confidence=getattr( + args, + "geometry_cross_view_outer_positive_confidence", + 0.70, + ), + geometry_cross_view_outer_positive_margin=getattr( + args, "geometry_cross_view_outer_positive_margin", 0.20 + ), + geometry_cross_view_outer_negative_confidence=getattr( + args, + "geometry_cross_view_outer_negative_confidence", + 0.70, + ), + geometry_cross_view_outer_negative_margin=getattr( + args, "geometry_cross_view_outer_negative_margin", 0.20 + ), + geometry_cross_view_outer_background_max_coverage=getattr( + args, + "geometry_cross_view_outer_background_max_coverage", + 0.25, + ), + geometry_cross_view_outer_min_views=getattr( + args, "geometry_cross_view_outer_min_views", 2 + ), + outer_uv_occupancy=getattr( + args, "outer_uv_occupancy_routing", False + ), + outer_uv_occupancy_blend_weight=getattr( + args, "outer_uv_occupancy_blend_weight", 0.30 + ), + outer_uv_occupancy_gate_threshold=getattr( + args, "outer_uv_occupancy_gate_threshold", 0.15 + ), + outer_uv_occupancy_rescue_threshold=getattr( + args, "outer_uv_occupancy_rescue_threshold", 0.70 + ), + outer_uv_occupancy_rescue_route_threshold=getattr( + args, + "outer_uv_occupancy_rescue_route_threshold", + 0.30, + ), + outer_uv_component_routing=getattr( + args, "outer_uv_component_routing", False + ), + outer_uv_component_seed_threshold=getattr( + args, "outer_uv_component_seed_threshold", 0.80 + ), + outer_uv_component_grow_threshold=getattr( + args, "outer_uv_component_grow_threshold", 0.50 + ), + outer_uv_component_min_size=getattr( + args, "outer_uv_component_min_size", 2 + ), + observed_foreground=None, + background_color_tolerance=getattr( + args, "background_color_tolerance", 0.25 + ), + color_background_tolerance=getattr( + args, "color_background_tolerance", 8.0 / 255.0 + ), + color_foreground_inset=getattr( + args, "color_foreground_inset", 1 + ), + reject_semantic_fallback=not args.allow_semantic_fallback, + return_details=True, + ) + gt_conditioning = splat_deterministic_targets_to_uv_conditioning( + rendered, + targets, + renderer=renderer, + views=views, + group_size=view_count, + bg_color=args.bg_color, + ) + else: + routing_details = None + pred_conditioning = splat_predictions_to_uv_conditioning( + rendered, + outputs, + group_size=view_count, + fg_threshold=args.splat_fg_threshold, + bg_color=args.bg_color, + ) + gt_conditioning = splat_targets_to_uv_conditioning( + rendered, + targets, + group_size=view_count, + bg_color=args.bg_color, + ) + + count = min(max_items, pred_conditioning.shape[0]) + pred_conditioning = pred_conditioning[:count].detach().cpu() + gt_conditioning = gt_conditioning[:count].detach().cpu() + preview = torch.cat( + [ + pred_conditioning[:, 0:3], + pred_conditioning[:, 5:8], + gt_conditioning[:, 0:3], + gt_conditioning[:, 5:8], + ], + dim=0, + ) + save_image(preview.clamp(0.0, 1.0), output_path, nrow=count) + if "outer_uv_occupancy_logits" in outputs: + _, outer_part_masks = build_part_layer_masks() + outer_valid = outer_part_masks[:, 0].bool().any(dim=0) + outer_valid = outer_valid.to( + device=outputs["outer_uv_occupancy_logits"].device + ).view(1, 1, UV_SIZE, UV_SIZE) + predicted_occupancy = ( + torch.sigmoid(outputs["outer_uv_occupancy_logits"][:count].float()) + * outer_valid + ) + target_occupancy = ( + (batch["uv"][:count, 3:4] > args.target_alpha_threshold).float() + * outer_valid + ) + occupancy_preview = torch.cat( + [predicted_occupancy, target_occupancy], dim=0 + ) + occupancy_path = output_path.with_name( + f"{output_path.stem}_outer_occupancy{output_path.suffix}" + ) + save_image( + occupancy_preview.detach().cpu(), + occupancy_path, + nrow=count, + ) + + debug_count = min(count * view_count, rendered.shape[0]) + if routing_details is not None: + debug_outputs = {key: value[:debug_count] for key, value in routing_details["outputs"].items()} + rendered_debug = routing_details["rendered"][:debug_count] + targets_debug = { + key: value[:debug_count] + for key, value in canonicalize_dense_targets(targets).items() + } + routing_debug = { + key: value[:debug_count] + for key, value in routing_details["routing"].items() + } + pred_fg = routing_debug["foreground"] + pred_layer_values = routing_debug["layer"] + pred_uv = flat_uv_to_uv01(routing_debug["flat_uv"], rendered_debug.dtype) + else: + debug_outputs = {key: value[:debug_count] for key, value in outputs.items()} + rendered_debug = rendered[:debug_count] + targets_debug = {key: value[:debug_count] for key, value in targets.items()} + pred_fg = torch.sigmoid(debug_outputs["foreground"])[:, 0] > args.splat_fg_threshold + pred_layer_values = debug_outputs["layer"].argmax(dim=1) + pred_uv = prediction_uv01(debug_outputs) + gt_fg = targets_debug["foreground"][:, 0] > 0.5 + raw_fg = torch.sigmoid(debug_outputs["foreground"])[:, 0] > args.splat_fg_threshold + pred_part_values = ( + debug_outputs["part"].argmax(dim=1) + if "part" in debug_outputs + else routing_debug["part"] + ) + pred_part = torch.where(pred_fg, pred_part_values, torch.full_like(pred_part_values, IGNORE_INDEX)) + pred_layer = torch.where( + pred_fg, + pred_layer_values, + torch.full_like(debug_outputs["layer"].argmax(dim=1), IGNORE_INDEX), + ) + pred_face_values = ( + debug_outputs["face"].argmax(dim=1) + if "face" in debug_outputs + else routing_debug["face"] + ) + pred_face = torch.where( + pred_fg, + pred_face_values, + torch.full_like(pred_face_values, IGNORE_INDEX), + ) + gt_face = torch.where( + gt_fg, + targets_debug["face"], + torch.full_like(targets_debug["face"], IGNORE_INDEX), + ) + debug_images = [ + rendered_debug[:, :3], + colorize_foreground(pred_fg, args.bg_color, rendered_debug), + colorize_foreground(gt_fg, args.bg_color, rendered_debug), + colorize_labels(pred_part, PART_PALETTE, args.bg_color, rendered_debug), + colorize_labels(targets_debug["part"], PART_PALETTE, args.bg_color, rendered_debug), + colorize_labels(pred_layer, LAYER_PALETTE, args.bg_color, rendered_debug), + colorize_labels(targets_debug["layer"], LAYER_PALETTE, args.bg_color, rendered_debug), + colorize_labels(pred_face, FACE_PALETTE, args.bg_color, rendered_debug), + colorize_labels(gt_face, FACE_PALETTE, args.bg_color, rendered_debug), + colorize_labels( + ( + torch.where( + pred_fg, + debug_outputs["layer_face"].argmax(dim=1), + torch.full_like(debug_outputs["layer_face"].argmax(dim=1), IGNORE_INDEX), + ) + if "layer_face" in debug_outputs + else combine_layer_face(pred_layer, pred_face) + ), + LAYER_FACE_PALETTE, + args.bg_color, + rendered_debug, + ), + colorize_labels( + combine_layer_face(targets_debug["layer"], gt_face), + LAYER_FACE_PALETTE, + args.bg_color, + rendered_debug, + ), + ] + if "route_role" in targets_debug and debug_outputs["layer"].shape[1] == 3: + pred_route_role = torch.where( + raw_fg, + debug_outputs["layer"].argmax(dim=1), + torch.full_like(debug_outputs["layer"].argmax(dim=1), IGNORE_INDEX), + ) + debug_images.extend( + [ + colorize_labels(pred_route_role, ROUTE_ROLE_PALETTE, args.bg_color, rendered_debug), + colorize_labels(targets_debug["route_role"], ROUTE_ROLE_PALETTE, args.bg_color, rendered_debug), + ] + ) + if "route_role_prior" in debug_outputs: + prior_route_role = torch.where( + gt_fg, + debug_outputs["route_role_prior"].argmax(dim=1), + torch.full_like(targets_debug["route_role"], IGNORE_INDEX), + ) + debug_images.append( + colorize_labels( + prior_route_role, + ROUTE_ROLE_PALETTE, + args.bg_color, + rendered_debug, + ) + ) + if routing_details is not None: + geometry_debug = build_geometry_grid_debug( + renderer, + views, + rendered_debug.shape[0], + rendered_debug, + bg_color=args.bg_color, + ) + inner_fill, outer_fill = fill_geometry_grid_debug( + rendered_debug, + pred_fg, + pred_layer_values, + geometry_debug, + bg_color=args.bg_color, + ) + inner_overlay, outer_overlay = overlay_geometry_grid_debug( + rendered_debug, + geometry_debug, + ) + inner_routed_overlay, outer_routed_overlay = overlay_geometry_grid_debug( + rendered_debug, + geometry_debug, + base_images=(inner_fill, outer_fill), + ) + debug_images.extend( + [ + inner_overlay, + outer_overlay, + inner_routed_overlay, + outer_routed_overlay, + geometry_debug[0], + geometry_debug[1], + inner_fill, + outer_fill, + ] + ) + if routing_details is not None: + pred_surface = torch.where( + pred_fg, + routing_debug["surface"], + torch.full_like(routing_debug["surface"], IGNORE_INDEX), + ) + debug_images.extend( + [ + colorize_surface(pred_surface, args.bg_color, rendered_debug), + colorize_surface(targets_debug["surface"], args.bg_color, rendered_debug), + ] + ) + debug_images.extend( + [ + colorize_uv(pred_uv, pred_fg, args.bg_color), + colorize_uv(targets_debug["uv"], gt_fg, args.bg_color), + ] + ) + debug_preview = torch.cat(debug_images, dim=0) + debug_path = output_path.with_name(f"{output_path.stem}_debug{output_path.suffix}") + save_image(debug_preview.clamp(0.0, 1.0).detach().cpu(), debug_path, nrow=view_count) + + +def save_checkpoint( + path, + model, + optimizer, + scaler, + epoch, + args, + metrics, + best_metric=None, + train_generator=None, + val_generator=None, +): + checkpoint = { + "epoch": epoch, + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "scaler": scaler.state_dict() if scaler is not None else None, + "reproducibility_state": capture_reproducibility_state( + train_generator, + val_generator, + ), + "args": vars(args), + "metrics": metrics, + "best_metric": best_metric, + "model_config": { + "input_channels": 4, + "base_channels": args.base_channels, + "uv_size": UV_SIZE, + "uv_classification": model.uv_classification, + "view_classes": model.view_classes, + "parser_mode": args.parser_mode, + "predict_affine": model.predict_affine, + "affine_translation_scale": model.affine_translation_scale, + "affine_scale_range": model.affine_scale_range, + "surface_classes": model.surface_classes, + "layer_classes": model.layer_classes, + "route_role_classes": model.layer_classes if model.geometry_only else 0, + "layer_face_classes": 0 if model.geometry_only else 12, + "geometry_only": model.geometry_only, + "feature_dropout": model.feature_dropout_probability, + "semantic_backbone": args.semantic_backbone, + "semantic_model": ( + args.tipsv2_model + if args.semantic_backbone == "tipsv2" + else args.siglip_model + if args.semantic_backbone == "siglip2" + else None + ), + "siglip_model": args.siglip_model if args.semantic_backbone == "siglip2" else None, + "tipsv2_model": args.tipsv2_model if args.semantic_backbone == "tipsv2" else None, + "semantic_feature_dim": model.semantic_feature_dim, + "semantic_channels": model.semantic_channels, + "semantic_attention_heads": model.semantic_attention_heads, + "semantic_layers": model.semantic_layers, + "semantic_dropout": model.semantic_dropout, + "semantic_spatial_feature_dim": model.semantic_spatial_feature_dim, + "semantic_spatial_channels": model.semantic_spatial_channels, + "predict_confidence": model.predict_confidence, + "route_role_spatial_prior": model.route_role_spatial_prior, + "route_prior_height": model.route_prior_height, + "route_prior_width": model.route_prior_width, + "route_prior_logit_cap": model.route_prior_logit_cap, + "route_prior_dropout": model.route_prior_dropout, + "predict_outer_uv_occupancy": ( + model.predict_outer_uv_occupancy + ), + "outer_uv_feature_channels": model.outer_uv_feature_channels, + "outer_uv_topology_channels": model.outer_uv_topology_channels, + "outer_uv_topology_layers": model.outer_uv_topology_layers, + "outer_uv_topology_dropout": model.outer_uv_topology_dropout, + "outer_uv_route_evidence_dropout": ( + model.outer_uv_route_evidence_dropout + ), + "arm_model": "steve", + }, + } + path = Path(path) + temporary_path = path.with_name(f".{path.name}.tmp") + torch.save(checkpoint, temporary_path) + os.replace(temporary_path, path) + + +def build_arg_parser(): + parser = argparse.ArgumentParser(description="Train a dense render-pixel to Minecraft UV parser.") + parser.add_argument("--data_dir", default="../skins") + parser.add_argument("--output_dir", default="runs/dense_uv_parser") + parser.add_argument("--resume", default=None, help="Checkpoint to resume; --epochs remains the final epoch number.") + parser.add_argument("--mappings_dir", default=None) + parser.add_argument("--views", default="walk_front_both_layer_ortho,walk_back_both_layer_ortho") + parser.add_argument( + "--privileged_views", + default="", + help=( + "Training-only renderer views arranged in complete groups matching " + "the primary --views roles. They are never used for validation or " + "checkpoint inference." + ), + ) + parser.add_argument( + "--parser_mode", + choices=["geometry_fit", "global_affine", "dense"], + default="geometry_fit", + help="geometry_fit learns alignment, route role, and exact renderer surface-slot routing.", + ) + parser.add_argument("--max_samples", type=int, default=180000) + parser.add_argument("--val_split", type=float, default=0.1) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--base_channels", type=int, default=32) + parser.add_argument("--feature_dropout", type=float, default=0.10) + parser.add_argument( + "--route_role_spatial_prior", + dest="route_role_spatial_prior", + action="store_true", + default=True, + help=( + "Learn a bounded per-view canonical route-role prior for common " + "fixed-view patterns such as bangs, hat brims, and back hair." + ), + ) + parser.add_argument( + "--no_route_role_spatial_prior", + dest="route_role_spatial_prior", + action="store_false", + ) + parser.add_argument("--route_prior_height", type=int, default=32) + parser.add_argument("--route_prior_width", type=int, default=16) + parser.add_argument("--route_prior_logit_cap", type=float, default=1.5) + parser.add_argument("--route_prior_dropout", type=float, default=0.10) + parser.add_argument( + "--semantic_backbone", + choices=["none", "siglip2", "tipsv2"], + default="none", + help=( + "Frozen semantic backbone. SigLIP2 can use cached global and spatial " + "features; TIPSv2 extracts both online." + ), + ) + parser.add_argument("--siglip_model", default="google/siglip2-base-patch16-224") + parser.add_argument("--siglip_cache_dir", default=None) + parser.add_argument("--siglip_cache_require_spatial", action="store_true") + parser.add_argument("--siglip_local_files_only", action="store_true") + parser.add_argument("--tipsv2_model", default="google/tipsv2-b14") + parser.add_argument("--tipsv2_local_files_only", action="store_true") + parser.add_argument("--semantic_channels", type=int, default=128) + parser.add_argument("--semantic_attention_heads", type=int, default=4) + parser.add_argument("--semantic_layers", type=int, default=1) + parser.add_argument("--semantic_dropout", type=float, default=0.05) + parser.add_argument("--semantic_spatial_channels", type=int, default=64) + parser.add_argument("--semantic_runtime_batch_size", type=int, default=32) + parser.add_argument( + "--predict_outer_uv_occupancy", + dest="predict_outer_uv_occupancy", + action="store_true", + default=False, + help=( + "Predict grouped 64x64 outer-layer occupancy as a semantic prior " + "for inner/outer routing." + ), + ) + parser.add_argument( + "--no_predict_outer_uv_occupancy", + dest="predict_outer_uv_occupancy", + action="store_false", + ) + parser.add_argument("--outer_uv_feature_channels", type=int, default=32) + parser.add_argument("--outer_uv_topology_channels", type=int, default=64) + parser.add_argument("--outer_uv_topology_layers", type=int, default=3) + parser.add_argument("--outer_uv_topology_dropout", type=float, default=0.05) + parser.add_argument( + "--outer_uv_route_evidence_dropout", type=float, default=1.0 + ) + parser.add_argument("--batch_size", type=int, default=8) + parser.add_argument( + "--gradient_accumulation_steps", type=int, default=1 + ) + parser.add_argument("--num_workers", type=int, default=8) + parser.add_argument("--prefetch_factor", type=int, default=4) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument( + "--lr_schedule", + choices=["constant", "cosine"], + default="cosine", + help="Absolute-epoch schedule; cosine is resume-safe and reduces late-epoch drift.", + ) + parser.add_argument("--min_lr_ratio", type=float, default=0.05) + parser.add_argument("--weight_decay", type=float, default=1e-4) + parser.add_argument("--device", default="auto") + parser.add_argument("--mixed_precision", choices=["no", "fp16", "bf16"], default="bf16") + parser.add_argument("--matmul_precision", choices=["highest", "high", "medium"], default="high") + parser.add_argument( + "--reproducible", + dest="reproducible", + action="store_true", + default=True, + help=( + "Seed all RNG streams, use seeded DataLoaders, and request " + "deterministic kernels where PyTorch supports them." + ), + ) + parser.add_argument( + "--no_reproducible", + dest="reproducible", + action="store_false", + ) + parser.add_argument( + "--strict_determinism", + action="store_true", + default=False, + help=( + "Fail instead of warning when CUDA has no deterministic " + "implementation for an operation." + ), + ) + parser.add_argument( + "--cudnn_benchmark", + dest="cudnn_benchmark", + action="store_true", + default=False, + ) + parser.add_argument("--no_cudnn_benchmark", dest="cudnn_benchmark", action="store_false") + parser.add_argument("--target_alpha_threshold", type=float, default=0.5) + parser.add_argument("--splat_fg_threshold", type=float, default=0.5) + parser.add_argument("--affine_refine", dest="affine_refine", action="store_true", default=False) + parser.add_argument("--no_affine_refine", dest="affine_refine", action="store_false") + parser.add_argument("--affine_refine_translation_px", type=float, default=0.0) + parser.add_argument("--affine_refine_scale", type=float, default=0.0) + parser.add_argument("--route_confidence_threshold", type=float, default=0.0) + parser.add_argument("--route_margin_threshold", type=float, default=0.0) + parser.add_argument( + "--background_color_tolerance", type=float, default=0.25 + ) + parser.add_argument( + "--color_background_tolerance", type=float, default=8.0 / 255.0 + ) + parser.add_argument("--color_foreground_inset", type=int, default=1) + parser.add_argument("--outer_route_confidence_threshold", type=float, default=0.80) + parser.add_argument("--outer_route_margin_threshold", type=float, default=0.55) + parser.add_argument("--outer_uv_min_coverage", type=float, default=0.25) + parser.add_argument("--outer_uv_min_source_pixels", type=int, default=15) + parser.add_argument( + "--outer_silhouette_consistency", + dest="outer_silhouette_consistency", + action="store_true", + default=True, + ) + parser.add_argument( + "--no_outer_silhouette_consistency", + dest="outer_silhouette_consistency", + action="store_false", + ) + parser.add_argument( + "--outer_silhouette_min_coverage", type=float, default=0.50 + ) + parser.add_argument("--outer_silhouette_dilation", type=int, default=0) + parser.add_argument("--outer_silhouette_min_pixels", type=int, default=4) + parser.add_argument( + "--outer_geometry_rescue", + dest="outer_geometry_rescue", + action="store_true", + default=True, + ) + parser.add_argument( + "--no_outer_geometry_rescue", + dest="outer_geometry_rescue", + action="store_false", + ) + parser.add_argument("--outer_rescue_confidence_threshold", type=float, default=0.60) + parser.add_argument("--outer_rescue_margin_threshold", type=float, default=0.25) + parser.add_argument("--outer_rescue_min_coverage", type=float, default=0.10) + parser.add_argument( + "--geometry_route_texel_consensus", + dest="geometry_route_texel_consensus", + action="store_true", + default=True, + help="Let fixed projected UV cells override local semantic route-role predictions.", + ) + parser.add_argument( + "--no_geometry_route_texel_consensus", + dest="geometry_route_texel_consensus", + action="store_false", + ) + parser.add_argument( + "--geometry_route_texel_consensus_weight", type=float, default=0.60 + ) + parser.add_argument( + "--geometry_route_preserve_outer_confidence", type=float, default=0.80 + ) + parser.add_argument( + "--geometry_route_preserve_outer_margin", type=float, default=0.35 + ) + parser.add_argument( + "--geometry_route_consensus_outer_confidence", type=float, default=0.70 + ) + parser.add_argument( + "--geometry_route_consensus_outer_margin", type=float, default=0.20 + ) + parser.add_argument( + "--geometry_cross_view_outer_consistency", + dest="geometry_cross_view_outer_consistency", + action="store_true", + default=False, + help=( + "Reject an outer route when another view provides strong " + "background/inner evidence for the same outer UV texel." + ), + ) + parser.add_argument( + "--no_geometry_cross_view_outer_consistency", + dest="geometry_cross_view_outer_consistency", + action="store_false", + ) + parser.add_argument( + "--geometry_cross_view_outer_weight", type=float, default=0.50 + ) + parser.add_argument( + "--geometry_cross_view_outer_positive_confidence", + type=float, + default=0.70, + ) + parser.add_argument( + "--geometry_cross_view_outer_positive_margin", + type=float, + default=0.20, + ) + parser.add_argument( + "--geometry_cross_view_outer_negative_confidence", + type=float, + default=0.70, + ) + parser.add_argument( + "--geometry_cross_view_outer_negative_margin", + type=float, + default=0.20, + ) + parser.add_argument( + "--geometry_cross_view_outer_background_max_coverage", + type=float, + default=0.25, + ) + parser.add_argument( + "--geometry_cross_view_outer_min_views", type=int, default=2 + ) + parser.add_argument( + "--outer_uv_occupancy_blend_weight", type=float, default=0.0 + ) + parser.add_argument( + "--outer_uv_occupancy_routing", + dest="outer_uv_occupancy_routing", + action="store_true", + default=False, + help=( + "Experimental: let the auxiliary occupancy head alter hard " + "routing and checkpoint selection." + ), + ) + parser.add_argument( + "--no_outer_uv_occupancy_routing", + dest="outer_uv_occupancy_routing", + action="store_false", + ) + parser.add_argument( + "--outer_uv_occupancy_gate_threshold", type=float, default=0.15 + ) + parser.add_argument( + "--outer_uv_occupancy_rescue_threshold", type=float, default=0.70 + ) + parser.add_argument( + "--outer_uv_occupancy_rescue_route_threshold", + type=float, + default=0.30, + ) + parser.add_argument( + "--outer_uv_component_routing", + dest="outer_uv_component_routing", + action="store_true", + default=False, + ) + parser.add_argument( + "--no_outer_uv_component_routing", + dest="outer_uv_component_routing", + action="store_false", + ) + parser.add_argument( + "--outer_uv_component_seed_threshold", type=float, default=0.80 + ) + parser.add_argument( + "--outer_uv_component_grow_threshold", type=float, default=0.50 + ) + parser.add_argument("--outer_uv_component_min_size", type=int, default=2) + parser.add_argument( + "--splat_color_aggregation", + choices=SPLAT_COLOR_AGGREGATIONS, + default="grid_mode", + ) + parser.add_argument("--allow_semantic_fallback", action="store_true") + parser.add_argument("--background_augment", dest="background_augment", action="store_true", default=True) + parser.add_argument("--no_background_augment", dest="background_augment", action="store_false") + parser.add_argument("--background_augment_prob", type=float, default=0.9) + parser.add_argument("--semantic_gate", dest="semantic_gate", action="store_true", default=True) + parser.add_argument("--no_semantic_gate", dest="semantic_gate", action="store_false") + parser.add_argument("--lambda_foreground", type=float, default=1.0) + parser.add_argument("--lambda_layer", type=float, default=1.0) + parser.add_argument("--lambda_part", type=float, default=0.5) + parser.add_argument("--lambda_face", type=float, default=0.5) + parser.add_argument("--lambda_layer_face", type=float, default=1.0) + parser.add_argument("--lambda_uv", type=float, default=0.25) + parser.add_argument("--lambda_uv_class", type=float, default=1.0) + parser.add_argument("--lambda_affine", type=float, default=1.0) + parser.add_argument("--lambda_surface", type=float, default=1.0) + parser.add_argument("--lambda_outer_false_positive", type=float, default=1.0) + parser.add_argument("--lambda_outer_false_negative", type=float, default=0.75) + parser.add_argument("--lambda_route_confidence", type=float, default=0.25) + parser.add_argument("--lambda_primary_route_swap", type=float, default=1.0) + parser.add_argument( + "--lambda_route_texel_consistency", type=float, default=0.25 + ) + parser.add_argument( + "--lambda_route_texel_supervision", + type=float, + default=0.0, + help="Weight for center-weighted, cross-view UV-texel route supervision.", + ) + parser.add_argument( + "--lambda_cross_view_outer_visibility", + type=float, + default=0.0, + help=( + "Weight for direct outer-alpha supervision over the union of " + "visible candidates and consistency over shared-view candidates." + ), + ) + parser.add_argument( + "--cross_view_outer_consistency_loss_weight", + type=float, + default=0.25, + help="Within-loss penalty for disagreement across shared views.", + ) + parser.add_argument( + "--lambda_privileged_view_distillation", + type=float, + default=0.0, + help=( + "Weight for UV-space distillation from training-only views into " + "the primary inference views." + ), + ) + parser.add_argument( + "--privileged_view_teacher_confidence", + type=float, + default=0.70, + help=( + "Minimum calibrated confidence for a privileged UV texel to act " + "as a stop-gradient teacher." + ), + ) + parser.add_argument( + "--outer_visibility_hard_negative_fraction", + type=float, + default=0.20, + help=( + "Fraction of visible transparent outer candidates retained as " + "high-confidence hard negatives." + ), + ) + parser.add_argument( + "--outer_visibility_hard_negative_weight", + type=float, + default=0.75, + help="Within-loss weight for mined outer false-positive candidates.", + ) + parser.add_argument( + "--route_texel_center_power", + type=float, + default=2.0, + help="Power applied to renderer texel-center quality in route supervision.", + ) + parser.add_argument( + "--lambda_route_prior_regularization", type=float, default=0.001 + ) + parser.add_argument("--lambda_semantic_presence", type=float, default=0.25) + parser.add_argument("--lambda_semantic_coverage", type=float, default=0.25) + parser.add_argument("--lambda_outer_uv_occupancy", type=float, default=0.0) + parser.add_argument( + "--outer_uv_occupancy_dice_weight", type=float, default=0.50 + ) + parser.add_argument( + "--outer_uv_occupancy_positive_balance", + type=float, + default=0.60, + ) + parser.add_argument( + "--outer_hard_positive_fraction", type=float, default=0.25 + ) + parser.add_argument( + "--outer_hard_positive_weight", type=float, default=0.50 + ) + parser.add_argument( + "--outer_hard_negative_fraction", type=float, default=0.02 + ) + parser.add_argument( + "--outer_hard_negative_weight", type=float, default=0.25 + ) + parser.add_argument( + "--lambda_outer_component_recall", type=float, default=0.0 + ) + parser.add_argument( + "--lambda_outer_component_false_positive", + type=float, + default=0.0, + ) + parser.add_argument("--lambda_outer_topology", type=float, default=0.0) + parser.add_argument( + "--lambda_outer_negative_topology", type=float, default=0.0 + ) + parser.add_argument( + "--lambda_route_occupancy_agreement", type=float, default=0.0 + ) + parser.add_argument( + "--outer_occupancy_agreement_warmup_fraction", + type=float, + default=0.25, + ) + parser.add_argument( + "--outer_occupancy_agreement_confidence_threshold", + type=float, + default=0.80, + ) + parser.add_argument("--outer_false_positive_gamma", type=float, default=3.0) + parser.add_argument("--outer_false_negative_gamma", type=float, default=2.0) + parser.add_argument("--primary_route_swap_gamma", type=float, default=2.0) + parser.add_argument("--route_prior_tv_weight", type=float, default=1.0) + parser.add_argument("--route_class_weight_floor", type=float, default=0.75) + parser.add_argument("--route_outer_class_weight_cap", type=float, default=0.90) + parser.add_argument("--lambda_soft_uv_rgb", type=float, default=0.25) + parser.add_argument("--lambda_soft_uv_alpha", type=float, default=0.35) + parser.add_argument("--lambda_soft_uv_inner_recall", type=float, default=0.50) + parser.add_argument("--lambda_soft_uv_outer_recall", type=float, default=0.50) + parser.add_argument("--soft_uv_recall_hard_fraction", type=float, default=0.10) + parser.add_argument("--soft_uv_recall_hard_weight", type=float, default=0.50) + parser.add_argument("--lambda_render_rgb", type=float, default=0.20) + parser.add_argument("--lambda_render_alpha", type=float, default=0.25) + parser.add_argument( + "--lambda_outer_projection_false_positive", + type=float, + default=0.0, + ) + parser.add_argument( + "--lambda_outer_projection_false_negative", + type=float, + default=0.0, + ) + parser.add_argument( + "--lambda_outer_projection_dice", type=float, default=0.0 + ) + parser.add_argument( + "--lambda_outer_projected_area", type=float, default=0.0 + ) + parser.add_argument("--outer_selection_precision_weight", type=float, default=1.5) + parser.add_argument("--outer_selection_recall_weight", type=float, default=0.5) + parser.add_argument("--outer_selection_iou_weight", type=float, default=0.5) + parser.add_argument("--inner_selection_recall_weight", type=float, default=0.5) + parser.add_argument("--hard_rgb_selection_weight", type=float, default=1.0) + parser.add_argument( + "--outer_projection_fp_selection_weight", + type=float, + default=0.0, + ) + parser.add_argument( + "--outer_projection_area_selection_weight", + type=float, + default=0.0, + ) + parser.add_argument( + "--render_softmax_temperature", + type=float, + default=1.0, + help="Temperature for the differentiable route/surface probabilities.", + ) + parser.add_argument("--uv_classification", dest="uv_classification", action="store_true", default=True) + parser.add_argument("--no_uv_classification", dest="uv_classification", action="store_false") + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_every", type=int, default=50) + parser.add_argument("--save_every", type=int, default=1) + parser.add_argument("--preview_every", type=int, default=1) + parser.add_argument( + "--best_metric", + choices=[ + "loss_total", + "loss_geometry", + "loss_routing", + "loss_surface", + "loss_uv_class", + "loss_differentiable", + "loss_outer_selection", + "loss_hard_uv_selection", + "loss_hard_uv_color_selection", + ], + default="loss_hard_uv_color_selection", + ) + return parser + + +def main(): + args = build_arg_parser().parse_args() + args.bg_color = (128, 128, 128) + if args.strict_determinism and not args.reproducible: + raise ValueError("--strict_determinism requires --reproducible.") + seed_everything(args.seed, reproducible=args.reproducible) + if not 0.0 <= args.feature_dropout < 1.0: + raise ValueError("--feature_dropout must be in [0, 1).") + if args.gradient_accumulation_steps < 1: + raise ValueError("--gradient_accumulation_steps must be positive.") + if args.route_prior_height < 1 or args.route_prior_width < 1: + raise ValueError("Route-prior dimensions must be positive.") + if args.route_prior_logit_cap <= 0: + raise ValueError("--route_prior_logit_cap must be positive.") + if not 0.0 <= args.route_prior_dropout < 1.0: + raise ValueError("--route_prior_dropout must be in [0, 1).") + if args.route_role_spatial_prior and args.parser_mode != "geometry_fit": + raise ValueError( + "--route_role_spatial_prior is supported only by geometry_fit." + ) + if ( + args.lambda_privileged_view_distillation > 0.0 + and not parse_views(args.privileged_views) + ): + raise ValueError( + "--lambda_privileged_view_distillation requires " + "--privileged_views." + ) + differentiable_weights = ( + args.lambda_soft_uv_rgb, + args.lambda_soft_uv_alpha, + args.lambda_soft_uv_inner_recall, + args.lambda_soft_uv_outer_recall, + args.lambda_render_rgb, + args.lambda_render_alpha, + args.lambda_outer_projection_false_positive, + args.lambda_outer_projection_false_negative, + args.lambda_outer_projection_dice, + args.lambda_outer_projected_area, + args.lambda_cross_view_outer_visibility, + args.lambda_privileged_view_distillation, + ) + if any(weight < 0 for weight in differentiable_weights): + raise ValueError("Differentiable parser loss weights must be non-negative.") + if min( + args.hard_rgb_selection_weight, + args.outer_projection_fp_selection_weight, + args.outer_projection_area_selection_weight, + ) < 0: + raise ValueError("Checkpoint selection weights must be non-negative.") + if not 0.0 <= args.background_color_tolerance <= 1.0: + raise ValueError("--background_color_tolerance must be in [0, 1].") + if not 0.0 <= args.color_background_tolerance <= 1.0: + raise ValueError("--color_background_tolerance must be in [0, 1].") + if args.color_foreground_inset < 0: + raise ValueError("--color_foreground_inset must be non-negative.") + if args.outer_uv_min_source_pixels < 1: + raise ValueError("--outer_uv_min_source_pixels must be positive.") + if args.outer_silhouette_dilation < 0: + raise ValueError("--outer_silhouette_dilation must be non-negative.") + if args.outer_silhouette_min_pixels < 1: + raise ValueError("--outer_silhouette_min_pixels must be positive.") + for name in ( + "outer_silhouette_min_coverage", + "geometry_route_texel_consensus_weight", + "geometry_route_preserve_outer_confidence", + "geometry_route_preserve_outer_margin", + "geometry_route_consensus_outer_confidence", + "geometry_route_consensus_outer_margin", + "geometry_cross_view_outer_weight", + "geometry_cross_view_outer_positive_confidence", + "geometry_cross_view_outer_positive_margin", + "geometry_cross_view_outer_negative_confidence", + "geometry_cross_view_outer_negative_margin", + "geometry_cross_view_outer_background_max_coverage", + "outer_uv_occupancy_blend_weight", + "outer_uv_occupancy_gate_threshold", + "outer_uv_occupancy_rescue_threshold", + "outer_uv_occupancy_rescue_route_threshold", + "outer_uv_component_seed_threshold", + "outer_uv_component_grow_threshold", + ): + if not 0.0 <= getattr(args, name) <= 1.0: + raise ValueError(f"--{name} must be in [0, 1].") + if args.lr <= 0: + raise ValueError("--lr must be positive.") + if not 0.0 <= args.min_lr_ratio <= 1.0: + raise ValueError("--min_lr_ratio must be in [0, 1].") + if not 0.0 < args.soft_uv_recall_hard_fraction <= 1.0: + raise ValueError("--soft_uv_recall_hard_fraction must be in (0, 1].") + if not 0.0 <= args.soft_uv_recall_hard_weight <= 1.0: + raise ValueError("--soft_uv_recall_hard_weight must be in [0, 1].") + if args.lambda_outer_false_positive < 0: + raise ValueError("--lambda_outer_false_positive must be non-negative.") + if args.lambda_outer_false_negative < 0: + raise ValueError("--lambda_outer_false_negative must be non-negative.") + if min( + args.lambda_route_confidence, + args.lambda_primary_route_swap, + args.lambda_route_texel_consistency, + args.lambda_route_texel_supervision, + args.lambda_cross_view_outer_visibility, + args.cross_view_outer_consistency_loss_weight, + args.lambda_route_prior_regularization, + args.lambda_semantic_presence, + args.lambda_semantic_coverage, + args.lambda_outer_uv_occupancy, + args.outer_uv_occupancy_dice_weight, + args.outer_hard_positive_weight, + args.outer_hard_negative_weight, + args.lambda_outer_component_recall, + args.lambda_outer_component_false_positive, + args.lambda_outer_topology, + args.lambda_outer_negative_topology, + args.lambda_route_occupancy_agreement, + ) < 0: + raise ValueError("Route, semantic, and confidence loss weights must be non-negative.") + if args.route_texel_center_power < 0: + raise ValueError("--route_texel_center_power must be non-negative.") + if args.geometry_cross_view_outer_min_views < 2: + raise ValueError( + "--geometry_cross_view_outer_min_views must be at least 2." + ) + if not 0.0 <= args.cross_view_outer_consistency_loss_weight <= 1.0: + raise ValueError( + "--cross_view_outer_consistency_loss_weight must be in [0, 1]." + ) + if not 0.0 <= args.privileged_view_teacher_confidence <= 1.0: + raise ValueError( + "--privileged_view_teacher_confidence must be in [0, 1]." + ) + if not 0.0 <= args.outer_visibility_hard_negative_fraction <= 1.0: + raise ValueError( + "--outer_visibility_hard_negative_fraction must be in [0, 1]." + ) + if args.outer_visibility_hard_negative_weight < 0.0: + raise ValueError( + "--outer_visibility_hard_negative_weight must be non-negative." + ) + if not 0.0 <= args.outer_hard_positive_fraction <= 1.0: + raise ValueError("--outer_hard_positive_fraction must be in [0, 1].") + if not 0.0 <= args.outer_uv_occupancy_positive_balance <= 1.0: + raise ValueError( + "--outer_uv_occupancy_positive_balance must be in [0, 1]." + ) + if not 0.0 <= args.outer_hard_negative_fraction <= 1.0: + raise ValueError("--outer_hard_negative_fraction must be in [0, 1].") + if not ( + 0.0 + <= args.outer_uv_component_grow_threshold + <= args.outer_uv_component_seed_threshold + <= 1.0 + ): + raise ValueError( + "Require outer component grow <= seed threshold in [0, 1]." + ) + if args.outer_uv_component_min_size < 1: + raise ValueError("--outer_uv_component_min_size must be positive.") + if not 0.0 <= args.outer_occupancy_agreement_warmup_fraction < 1.0: + raise ValueError( + "--outer_occupancy_agreement_warmup_fraction must be in [0, 1)." + ) + if not 0.5 <= args.outer_occupancy_agreement_confidence_threshold <= 1.0: + raise ValueError( + "--outer_occupancy_agreement_confidence_threshold must be in " + "[0.5, 1]." + ) + if ( + args.outer_uv_feature_channels < 1 + or args.outer_uv_topology_channels < 1 + or args.outer_uv_topology_layers < 1 + ): + raise ValueError("Outer UV topology dimensions must be positive.") + if not 0.0 <= args.outer_uv_topology_dropout < 1.0: + raise ValueError("--outer_uv_topology_dropout must be in [0, 1).") + if not 0.0 <= args.outer_uv_route_evidence_dropout <= 1.0: + raise ValueError( + "--outer_uv_route_evidence_dropout must be in [0, 1]." + ) + if ( + args.outer_uv_occupancy_routing + or args.outer_uv_component_routing + or args.lambda_outer_uv_occupancy > 0.0 + or args.lambda_outer_component_recall > 0.0 + or args.lambda_outer_component_false_positive > 0.0 + or args.lambda_outer_topology > 0.0 + or args.lambda_outer_negative_topology > 0.0 + or args.lambda_route_occupancy_agreement > 0.0 + ) and not args.predict_outer_uv_occupancy: + raise ValueError( + "Outer occupancy routing/losses require " + "--predict_outer_uv_occupancy." + ) + if ( + args.semantic_channels < 1 + or args.semantic_layers < 1 + or args.semantic_spatial_channels < 1 + or args.semantic_runtime_batch_size < 1 + ): + raise ValueError("Semantic channels and layers must be positive.") + if args.semantic_attention_heads < 1: + raise ValueError("--semantic_attention_heads must be positive.") + if args.semantic_channels % args.semantic_attention_heads != 0: + raise ValueError("Semantic channels must be divisible by attention heads.") + if not 0.0 <= args.semantic_dropout < 1.0: + raise ValueError("--semantic_dropout must be in [0, 1).") + if args.outer_false_positive_gamma < 0: + raise ValueError("--outer_false_positive_gamma must be non-negative.") + if args.outer_false_negative_gamma < 0: + raise ValueError("--outer_false_negative_gamma must be non-negative.") + if args.primary_route_swap_gamma < 0: + raise ValueError("--primary_route_swap_gamma must be non-negative.") + if args.route_prior_tv_weight < 0: + raise ValueError("--route_prior_tv_weight must be non-negative.") + if not 0 <= args.route_class_weight_floor <= 4.0: + raise ValueError("--route_class_weight_floor must be in [0, 4].") + if args.route_outer_class_weight_cap <= 0: + raise ValueError("--route_outer_class_weight_cap must be positive.") + if ( + args.outer_selection_precision_weight < 0 + or args.outer_selection_recall_weight < 0 + or args.outer_selection_iou_weight < 0 + or args.inner_selection_recall_weight < 0 + ): + raise ValueError("Outer checkpoint-selection weights must be non-negative.") + if args.render_softmax_temperature <= 0: + raise ValueError("--render_softmax_temperature must be positive.") + device = get_device(args.device) + configure_torch(args, device) + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "previews").mkdir(exist_ok=True) + + primary_views, privileged_views, training_views = ( + privileged_training_views(args) + ) + if len(set(training_views)) != len(training_views): + raise ValueError( + "Primary and privileged renderer views must be unique." + ) + if privileged_views and args.predict_outer_uv_occupancy: + raise ValueError( + "The experimental occupancy head does not support training-only " + "view groups. Keep PREDICT_OUTER_UV_OCCUPANCY=false." + ) + + dataset = SkinUVDataset( + data_dir=args.data_dir, + mappings_dir=args.mappings_dir, + views=args.views, + max_samples=args.max_samples, + ) + semantic_cache = None + runtime_semantic_backbone = None + semantic_feature_dim = 0 + semantic_spatial_feature_dim = 0 + semantic_masks = None + semantic_model_name = ( + args.tipsv2_model + if args.semantic_backbone == "tipsv2" + else args.siglip_model + ) + semantic_local_files_only = ( + args.tipsv2_local_files_only + if args.semantic_backbone == "tipsv2" + else args.siglip_local_files_only + ) + if args.semantic_backbone != "siglip2" and args.siglip_cache_dir: + raise ValueError("--siglip_cache_dir can only be used with --semantic_backbone siglip2.") + if args.semantic_backbone == "siglip2": + if args.siglip_cache_dir: + semantic_cache = SigLIPGlobalCache( + args.siglip_cache_dir, + expected_views=training_views, + expected_model=args.siglip_model, + expected_data_dir=args.data_dir, + require_spatial=args.siglip_cache_require_spatial, + ) + missing_semantics = [ + Path(path).name + for path in dataset.skin_paths + if Path(path).name not in semantic_cache.filename_to_index + ] + if missing_semantics: + raise ValueError( + f"SigLIP cache is missing {len(missing_semantics)} selected skins; " + f"first missing: {missing_semantics[0]}." + ) + semantic_feature_dim = int(semantic_cache.metadata["feature_dim"]) + if semantic_cache.has_spatial: + semantic_spatial_feature_dim = int( + semantic_cache.metadata["spatial_feature_dim"] + ) + else: + runtime_semantic_backbone = build_semantic_runtime( + args.semantic_backbone, + semantic_model_name, + device, + semantic_channels=args.semantic_channels, + local_files_only=semantic_local_files_only, + runtime_batch_size=args.semantic_runtime_batch_size, + ) + semantic_feature_dim = runtime_semantic_backbone.raw_feature_dim + semantic_spatial_feature_dim = int( + getattr( + runtime_semantic_backbone, + "raw_spatial_feature_dim", + 0, + ) + ) + elif args.semantic_backbone == "tipsv2": + runtime_semantic_backbone = build_semantic_runtime( + args.semantic_backbone, + semantic_model_name, + device, + semantic_channels=args.semantic_channels, + local_files_only=semantic_local_files_only, + runtime_batch_size=args.semantic_runtime_batch_size, + ) + semantic_feature_dim = runtime_semantic_backbone.raw_feature_dim + semantic_spatial_feature_dim = ( + runtime_semantic_backbone.raw_spatial_feature_dim + ) + if args.semantic_backbone != "none" or args.predict_outer_uv_occupancy: + semantic_masks = tuple( + mask.to(device) for mask in build_part_layer_masks() + ) + val_count = int(len(dataset) * args.val_split) + train_count = len(dataset) - val_count + split_generator = torch.Generator().manual_seed(args.seed) + train_loader_generator = torch.Generator().manual_seed(args.seed + 1) + val_loader_generator = torch.Generator().manual_seed(args.seed + 2) + if val_count > 0: + train_dataset, val_dataset = random_split( + dataset, + [train_count, val_count], + generator=split_generator, + ) + else: + train_dataset, val_dataset = dataset, None + + loader_kwargs = { + "batch_size": args.batch_size, + "num_workers": args.num_workers, + "pin_memory": device.type == "cuda", + "worker_init_fn": seed_data_worker, + } + if args.num_workers > 0: + loader_kwargs["persistent_workers"] = not args.reproducible + loader_kwargs["prefetch_factor"] = args.prefetch_factor + train_loader = DataLoader( + train_dataset, + shuffle=True, + generator=train_loader_generator, + **loader_kwargs, + ) + val_loader = ( + DataLoader( + val_dataset, + shuffle=False, + generator=val_loader_generator, + **loader_kwargs, + ) + if val_dataset is not None + else None + ) + + renderer = DifferentiableRenderer( + mappings_dir=args.mappings_dir, + bg_color=tuple(channel / 255.0 for channel in args.bg_color), + ).to(device) + missing_views = [ + view for view in training_views if view not in renderer.views + ] + if missing_views: + raise ValueError(f"Unknown renderer views {missing_views}. Available views: {', '.join(renderer.views)}") + + affine_mode = args.parser_mode in ("geometry_fit", "global_affine") + geometry_only = args.parser_mode == "geometry_fit" + surface_classes = ( + surface_class_count(renderer, training_views) + if args.parser_mode in ("geometry_fit", "global_affine") + else 0 + ) + model = DenseUVParserNet( + base_channels=args.base_channels, + uv_size=UV_SIZE, + uv_classification=args.uv_classification and not geometry_only, + # Only the primary inference roles receive distinct embeddings. Every + # privileged group reuses the same ordered role ids. + view_classes=len(primary_views), + predict_affine=affine_mode, + affine_translation_scale=0.0, + affine_scale_range=0.0, + surface_classes=surface_classes, + geometry_only=geometry_only, + feature_dropout=args.feature_dropout, + semantic_feature_dim=semantic_feature_dim, + semantic_channels=args.semantic_channels, + semantic_attention_heads=args.semantic_attention_heads, + semantic_layers=args.semantic_layers, + semantic_dropout=args.semantic_dropout, + semantic_spatial_feature_dim=semantic_spatial_feature_dim, + semantic_spatial_channels=args.semantic_spatial_channels, + predict_confidence=args.semantic_backbone != "none", + route_role_spatial_prior=( + geometry_only and args.route_role_spatial_prior + ), + route_prior_height=args.route_prior_height, + route_prior_width=args.route_prior_width, + route_prior_logit_cap=args.route_prior_logit_cap, + route_prior_dropout=args.route_prior_dropout, + predict_outer_uv_occupancy=args.predict_outer_uv_occupancy, + outer_uv_feature_channels=args.outer_uv_feature_channels, + outer_uv_topology_channels=args.outer_uv_topology_channels, + outer_uv_topology_layers=args.outer_uv_topology_layers, + outer_uv_topology_dropout=args.outer_uv_topology_dropout, + outer_uv_route_evidence_dropout=( + args.outer_uv_route_evidence_dropout + ), + ).to(device) + if runtime_semantic_backbone is not None: + attach_semantic_runtime( + model, + args.semantic_backbone, + semantic_model_name, + device, + local_files_only=semantic_local_files_only, + backbone=runtime_semantic_backbone, + runtime_batch_size=args.semantic_runtime_batch_size, + ) + criterion = DenseUVParserLoss( + lambda_foreground=args.lambda_foreground, + lambda_layer=args.lambda_layer, + lambda_part=args.lambda_part, + lambda_face=args.lambda_face, + lambda_layer_face=args.lambda_layer_face, + lambda_uv=args.lambda_uv, + lambda_uv_class=args.lambda_uv_class, + lambda_affine=args.lambda_affine, + lambda_surface=args.lambda_surface, + lambda_outer_false_positive=args.lambda_outer_false_positive, + lambda_outer_false_negative=args.lambda_outer_false_negative, + lambda_route_confidence=args.lambda_route_confidence, + lambda_primary_route_swap=args.lambda_primary_route_swap, + lambda_route_texel_consistency=args.lambda_route_texel_consistency, + lambda_route_prior_regularization=( + args.lambda_route_prior_regularization + ), + outer_false_positive_gamma=args.outer_false_positive_gamma, + outer_false_negative_gamma=args.outer_false_negative_gamma, + primary_route_swap_gamma=args.primary_route_swap_gamma, + route_prior_tv_weight=args.route_prior_tv_weight, + route_class_weight_floor=args.route_class_weight_floor, + route_outer_class_weight_cap=args.route_outer_class_weight_cap, + uv_size=UV_SIZE, + use_uv=(not affine_mode) or model.uv_classification, + affine_translation_limit=model.affine_translation_limit if affine_mode else 1.0, + affine_log_scale_limit=model.affine_log_scale_limit if affine_mode else 1.0, + ).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + scaler = build_grad_scaler(device, args.mixed_precision) + + start_epoch = 1 + best_metric = float("inf") + if args.resume: + checkpoint = torch.load(args.resume, map_location=device) + checkpoint_seed = checkpoint.get("args", {}).get("seed") + if ( + args.reproducible + and checkpoint_seed is not None + and int(checkpoint_seed) != args.seed + ): + raise ValueError( + "Cannot change --seed while reproducibly resuming: " + f"checkpoint={checkpoint_seed}, requested={args.seed}." + ) + checkpoint_mode = checkpoint.get("model_config", {}).get( + "parser_mode", checkpoint.get("args", {}).get("parser_mode", "dense") + ) + if checkpoint_mode != args.parser_mode: + raise ValueError( + f"Cannot resume parser_mode={checkpoint_mode!r} as {args.parser_mode!r}. Start a new run." + ) + checkpoint_semantic_backbone = checkpoint.get("model_config", {}).get( + "semantic_backbone", "none" + ) + if checkpoint_semantic_backbone != args.semantic_backbone: + raise ValueError( + "Cannot change semantic backbone while resuming: " + f"checkpoint={checkpoint_semantic_backbone!r}, " + f"requested={args.semantic_backbone!r}. Start a new run." + ) + if args.semantic_backbone != "none": + checkpoint_config = checkpoint.get("model_config", {}) + checkpoint_semantic_model = checkpoint_config.get( + "semantic_model", + checkpoint_config.get( + "tipsv2_model" + if args.semantic_backbone == "tipsv2" + else "siglip_model", + checkpoint.get("args", {}).get( + "tipsv2_model" + if args.semantic_backbone == "tipsv2" + else "siglip_model" + ), + ), + ) + if checkpoint_semantic_model != semantic_model_name: + raise ValueError( + "Cannot change semantic model while resuming: " + f"checkpoint={checkpoint_semantic_model!r}, " + f"requested={semantic_model_name!r}. Start a new run." + ) + semantic_resume_fields = { + "semantic_feature_dim": model.semantic_feature_dim, + "semantic_channels": model.semantic_channels, + "semantic_attention_heads": model.semantic_attention_heads, + "semantic_layers": model.semantic_layers, + "semantic_spatial_feature_dim": ( + model.semantic_spatial_feature_dim + ), + "semantic_spatial_channels": model.semantic_spatial_channels, + } + mismatches = { + name: ( + checkpoint_config.get( + name, + 0 + if name == "semantic_spatial_feature_dim" + else expected + if name == "semantic_spatial_channels" + else None, + ), + expected, + ) + for name, expected in semantic_resume_fields.items() + if checkpoint_config.get( + name, + 0 + if name == "semantic_spatial_feature_dim" + else expected + if name == "semantic_spatial_channels" + else None, + ) + != expected + } + if mismatches: + raise ValueError( + "Cannot change semantic adapter shape while resuming: " + f"{mismatches}. Start a new run." + ) + checkpoint_layer_classes = checkpoint.get("model_config", {}).get("layer_classes", 2) + if geometry_only and checkpoint_layer_classes != model.layer_classes: + raise ValueError( + "This geometry checkpoint predates the secondary/backface route class. " + "Start a new parser run instead of resuming it." + ) + checkpoint_surface_classes = checkpoint.get("model_config", {}).get("surface_classes", 0) + if geometry_only and checkpoint_surface_classes != model.surface_classes: + raise ValueError( + "This geometry checkpoint predates exact surface-slot routing. " + "Start a new parser run instead of resuming it." + ) + checkpoint_config = checkpoint.get("model_config", {}) + checkpoint_outer_occupancy = checkpoint_config.get( + "predict_outer_uv_occupancy", + any( + key.startswith("outer_uv_occupancy_head.") + for key in checkpoint["model"] + ), + ) + if ( + checkpoint_outer_occupancy + != model.predict_outer_uv_occupancy + ): + raise ValueError( + "Cannot add or remove the outer UV occupancy head while " + "resuming. Start a new parser run." + ) + checkpoint_route_prior = checkpoint_config.get( + "route_role_spatial_prior", + "route_role_prior" in checkpoint["model"], + ) + if geometry_only and checkpoint_route_prior != model.route_role_spatial_prior: + raise ValueError( + "Cannot add or remove the learned route-role spatial prior while " + "resuming. Start a new parser run, or explicitly disable the prior " + "to continue a legacy checkpoint." + ) + if checkpoint_route_prior: + prior_shape = tuple(checkpoint["model"]["route_role_prior"].shape[-2:]) + expected_prior_shape = ( + model.route_prior_height, + model.route_prior_width, + ) + if prior_shape != expected_prior_shape: + raise ValueError( + "Route-prior grid mismatch while resuming: " + f"checkpoint={prior_shape}, requested={expected_prior_shape}." + ) + if args.parser_mode == "global_affine" and not any( + key.startswith("layer_face.") for key in checkpoint["model"] + ): + raise ValueError("This global-affine checkpoint predates the joint layer-face head.") + model.load_state_dict(checkpoint["model"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + for param_group in optimizer.param_groups: + param_group["lr"] = args.lr + param_group["weight_decay"] = args.weight_decay + if scaler is not None and checkpoint.get("scaler") is not None: + scaler.load_state_dict(checkpoint["scaler"]) + if args.reproducible: + restored_rng = restore_reproducibility_state( + checkpoint.get("reproducibility_state"), + train_loader_generator, + val_loader_generator, + ) + if not restored_rng: + print( + "warning: checkpoint has no reproducibility state; " + "resume will use the requested seed but cannot exactly " + "continue the old random sequence." + ) + + checkpoint_epoch = int(checkpoint.get("epoch", 0)) + start_epoch = checkpoint_epoch + 1 + checkpoint_args = checkpoint.get("args", {}) + checkpoint_best_metric = checkpoint_args.get("best_metric", "loss_total") + if checkpoint_best_metric != args.best_metric: + print( + "Resetting best-metric history because checkpoint selection changed: " + f"checkpoint={checkpoint_best_metric!r}, requested={args.best_metric!r}." + ) + best_metric = float("inf") + else: + best_metric = float(checkpoint.get("best_metric", float("inf"))) + if start_epoch > args.epochs: + raise ValueError( + f"Checkpoint is already at epoch {checkpoint_epoch}; set --epochs above {checkpoint_epoch}." + ) + print( + json.dumps( + { + "resume": str(args.resume), + "checkpoint_epoch": checkpoint_epoch, + "start_epoch": start_epoch, + "target_epoch": args.epochs, + "lr": args.lr, + "best_metric": args.best_metric, + "best_metric_value": best_metric, + }, + indent=2, + ) + ) + + metadata = { + "num_samples": len(dataset), + "train_samples": len(train_dataset), + "val_samples": len(val_dataset) if val_dataset is not None else 0, + "seed": args.seed, + "reproducible": args.reproducible, + "strict_determinism": args.strict_determinism, + "deterministic_algorithms": ( + torch.are_deterministic_algorithms_enabled() + ), + "cudnn_benchmark": ( + torch.backends.cudnn.benchmark + if device.type == "cuda" + else False + ), + "cublas_workspace_config": os.environ.get( + "CUBLAS_WORKSPACE_CONFIG" + ), + "views": primary_views, + "privileged_views": privileged_views, + "training_views": training_views, + "parameters": count_parameters(model), + "batch_size": args.batch_size, + "gradient_accumulation_steps": args.gradient_accumulation_steps, + "effective_skin_batch_size": ( + args.batch_size * args.gradient_accumulation_steps + ), + "feature_dropout": args.feature_dropout, + "semantic_backbone": args.semantic_backbone, + "semantic_model": ( + semantic_model_name if args.semantic_backbone != "none" else None + ), + "semantic_feature_dim": semantic_feature_dim, + "semantic_channels": args.semantic_channels, + "semantic_layers": args.semantic_layers, + "semantic_spatial_feature_dim": semantic_spatial_feature_dim, + "semantic_spatial_channels": args.semantic_spatial_channels, + "semantic_runtime_batch_size": args.semantic_runtime_batch_size, + "predict_confidence": model.predict_confidence, + "route_role_spatial_prior": model.route_role_spatial_prior, + "route_prior_height": model.route_prior_height, + "route_prior_width": model.route_prior_width, + "route_prior_logit_cap": model.route_prior_logit_cap, + "route_prior_dropout": model.route_prior_dropout, + "predict_outer_uv_occupancy": model.predict_outer_uv_occupancy, + "outer_uv_occupancy_supervision": ( + "visible_projected_texels" + if model.predict_outer_uv_occupancy + else "disabled" + ), + "outer_uv_feature_channels": model.outer_uv_feature_channels, + "outer_uv_topology_channels": model.outer_uv_topology_channels, + "outer_uv_topology_layers": model.outer_uv_topology_layers, + "outer_uv_route_evidence_dropout": ( + model.outer_uv_route_evidence_dropout + ), + "outer_uv_occupancy_routing": args.outer_uv_occupancy_routing, + "outer_uv_component_routing": args.outer_uv_component_routing, + "outer_uv_component_seed_threshold": ( + args.outer_uv_component_seed_threshold + ), + "outer_uv_component_grow_threshold": ( + args.outer_uv_component_grow_threshold + ), + "outer_uv_component_min_size": args.outer_uv_component_min_size, + "device": str(device), + "parser_mode": args.parser_mode, + "uv_classification": model.uv_classification, + "view_classes": model.view_classes, + "surface_classes": surface_classes, + "layer_classes": model.layer_classes, + "route_role_classes": model.layer_classes if geometry_only else 0, + "geometry_only": geometry_only, + "arm_model": "steve", + "best_metric": args.best_metric, + "base_learning_rate": args.lr, + "lr_schedule": args.lr_schedule, + "min_lr_ratio": args.min_lr_ratio, + "background_augment": args.background_augment, + "background_augment_prob": args.background_augment_prob, + "semantic_gate": args.semantic_gate, + "geometry_route_texel_consensus": args.geometry_route_texel_consensus, + "geometry_cross_view_outer_consistency": ( + args.geometry_cross_view_outer_consistency + ), + "geometry_cross_view_outer_weight": ( + args.geometry_cross_view_outer_weight + ), + "geometry_cross_view_outer_min_views": ( + args.geometry_cross_view_outer_min_views + ), + "outer_uv_min_source_pixels": args.outer_uv_min_source_pixels, + "outer_silhouette_consistency": args.outer_silhouette_consistency, + "outer_silhouette_min_coverage": ( + args.outer_silhouette_min_coverage + ), + "outer_silhouette_dilation": args.outer_silhouette_dilation, + "outer_silhouette_min_pixels": args.outer_silhouette_min_pixels, + "background_color_tolerance": args.background_color_tolerance, + "color_background_tolerance": args.color_background_tolerance, + "color_foreground_inset": args.color_foreground_inset, + "affine_refine": args.affine_refine, + "affine_refine_translation_px": args.affine_refine_translation_px, + "affine_refine_scale": args.affine_refine_scale, + "lambda_soft_uv_rgb": args.lambda_soft_uv_rgb, + "lambda_soft_uv_alpha": args.lambda_soft_uv_alpha, + "lambda_soft_uv_inner_recall": args.lambda_soft_uv_inner_recall, + "lambda_soft_uv_outer_recall": args.lambda_soft_uv_outer_recall, + "soft_uv_recall_hard_fraction": args.soft_uv_recall_hard_fraction, + "soft_uv_recall_hard_weight": args.soft_uv_recall_hard_weight, + "lambda_render_rgb": args.lambda_render_rgb, + "lambda_render_alpha": args.lambda_render_alpha, + "lambda_outer_false_positive": args.lambda_outer_false_positive, + "lambda_outer_false_negative": args.lambda_outer_false_negative, + "lambda_route_confidence": args.lambda_route_confidence, + "lambda_primary_route_swap": args.lambda_primary_route_swap, + "lambda_route_texel_consistency": args.lambda_route_texel_consistency, + "lambda_route_texel_supervision": ( + args.lambda_route_texel_supervision + ), + "lambda_cross_view_outer_visibility": ( + args.lambda_cross_view_outer_visibility + ), + "cross_view_outer_consistency_loss_weight": ( + args.cross_view_outer_consistency_loss_weight + ), + "lambda_privileged_view_distillation": ( + args.lambda_privileged_view_distillation + ), + "privileged_view_teacher_confidence": ( + args.privileged_view_teacher_confidence + ), + "outer_visibility_hard_negative_fraction": ( + args.outer_visibility_hard_negative_fraction + ), + "outer_visibility_hard_negative_weight": ( + args.outer_visibility_hard_negative_weight + ), + "route_texel_center_power": args.route_texel_center_power, + "lambda_route_prior_regularization": ( + args.lambda_route_prior_regularization + ), + "lambda_semantic_presence": args.lambda_semantic_presence, + "lambda_semantic_coverage": args.lambda_semantic_coverage, + "lambda_outer_uv_occupancy": args.lambda_outer_uv_occupancy, + "outer_uv_occupancy_dice_weight": ( + args.outer_uv_occupancy_dice_weight + ), + "outer_uv_occupancy_positive_balance": ( + args.outer_uv_occupancy_positive_balance + ), + "outer_hard_positive_fraction": args.outer_hard_positive_fraction, + "outer_hard_positive_weight": args.outer_hard_positive_weight, + "outer_hard_negative_fraction": args.outer_hard_negative_fraction, + "outer_hard_negative_weight": args.outer_hard_negative_weight, + "lambda_outer_component_recall": ( + args.lambda_outer_component_recall + ), + "lambda_outer_component_false_positive": ( + args.lambda_outer_component_false_positive + ), + "lambda_outer_topology": args.lambda_outer_topology, + "lambda_outer_negative_topology": ( + args.lambda_outer_negative_topology + ), + "lambda_route_occupancy_agreement": ( + args.lambda_route_occupancy_agreement + ), + "outer_occupancy_agreement_warmup_fraction": ( + args.outer_occupancy_agreement_warmup_fraction + ), + "outer_occupancy_agreement_confidence_threshold": ( + args.outer_occupancy_agreement_confidence_threshold + ), + "outer_false_positive_gamma": args.outer_false_positive_gamma, + "outer_false_negative_gamma": args.outer_false_negative_gamma, + "primary_route_swap_gamma": args.primary_route_swap_gamma, + "route_prior_tv_weight": args.route_prior_tv_weight, + "route_class_weight_floor": args.route_class_weight_floor, + "route_outer_class_weight_cap": args.route_outer_class_weight_cap, + "outer_selection_precision_weight": args.outer_selection_precision_weight, + "outer_selection_recall_weight": args.outer_selection_recall_weight, + "outer_selection_iou_weight": args.outer_selection_iou_weight, + "inner_selection_recall_weight": args.inner_selection_recall_weight, + "hard_rgb_selection_weight": args.hard_rgb_selection_weight, + "lambda_outer_projection_false_positive": ( + args.lambda_outer_projection_false_positive + ), + "lambda_outer_projection_false_negative": ( + args.lambda_outer_projection_false_negative + ), + "lambda_outer_projection_dice": args.lambda_outer_projection_dice, + "lambda_outer_projected_area": args.lambda_outer_projected_area, + "outer_projection_fp_selection_weight": ( + args.outer_projection_fp_selection_weight + ), + "outer_projection_area_selection_weight": ( + args.outer_projection_area_selection_weight + ), + "render_softmax_temperature": args.render_softmax_temperature, + } + with open(output_dir / "config.json", "w", encoding="utf-8") as handle: + json.dump({"args": vars(args), "metadata": metadata}, handle, indent=2) + print(json.dumps(metadata, indent=2)) + + for epoch in range(start_epoch, args.epochs + 1): + epoch_lr = learning_rate_for_epoch( + args.lr, + epoch, + args.epochs, + schedule=args.lr_schedule, + min_lr_ratio=args.min_lr_ratio, + ) + for param_group in optimizer.param_groups: + param_group["lr"] = epoch_lr + train_metrics = run_epoch( + model, + criterion, + renderer, + train_loader, + optimizer, + scaler, + device, + args.mixed_precision, + args, + train=True, + compute_hard_metrics=( + val_loader is None + and args.best_metric + in ("loss_hard_uv_selection", "loss_hard_uv_color_selection") + ), + semantic_cache=semantic_cache, + semantic_masks=semantic_masks, + ) + train_metrics["learning_rate"] = epoch_lr + metrics = {"train": train_metrics} + metric_source = train_metrics + if val_loader is not None: + with torch.no_grad(): + val_metrics = run_epoch( + model, + criterion, + renderer, + val_loader, + optimizer, + scaler, + device, + args.mixed_precision, + args, + train=False, + compute_hard_metrics=True, + semantic_cache=semantic_cache, + semantic_masks=semantic_masks, + ) + metrics["val"] = val_metrics + metric_source = val_metrics + + if args.best_metric not in metric_source: + raise ValueError(f"Best metric {args.best_metric!r} is not available in epoch metrics.") + metric = metric_source[args.best_metric] + print(f"epoch={epoch} metrics={json.dumps(metrics, sort_keys=True)}") + + is_best = metric < best_metric + if is_best: + best_metric = metric + if epoch % args.preview_every == 0: + save_preview( + model, + renderer, + val_loader or train_loader, + device, + args, + output_dir / "previews" / f"epoch_{epoch:04d}.png", + semantic_cache=semantic_cache, + ) + if epoch % args.save_every == 0: + save_checkpoint( + output_dir / "latest.pt", + model, + optimizer, + scaler, + epoch, + args, + metrics, + best_metric=best_metric, + train_generator=train_loader_generator, + val_generator=val_loader_generator, + ) + if is_best: + save_checkpoint( + output_dir / "best.pt", + model, + optimizer, + scaler, + epoch, + args, + metrics, + best_metric=best_metric, + train_generator=train_loader_generator, + val_generator=val_loader_generator, + ) + + +if __name__ == "__main__": + main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train_accessories.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_accessories.py new file mode 100644 index 0000000000000000000000000000000000000000..583952a60435d2644374dfa79c4f42066ec86b84 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_accessories.py @@ -0,0 +1,311 @@ +"""v101 object-supervised adaptation of v61 with end-to-end UV validation.""" +import argparse +import hashlib +import inspect +import json +import math +import os +import shutil +import subprocess +import time +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from torchvision.utils import save_image +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet +from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime +from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset, load_skin +from SkingToolkit.dense_uv_parser.accessories import accessory_loss, head_bounds +from SkingToolkit.dense_uv_parser.accessory_data import AccessoryDataset, render_accessories +from SkingToolkit.dense_uv_parser.accessory_pipeline import run_pipeline, load_pipeline +from SkingToolkit.dense_uv_parser.semantic_generalization import appearance_augment, write_json +from SkingToolkit.dense_uv_parser.utils import build_static_surface_routing, build_dense_parser_batch +from SkingToolkit.dense_uv_parser.uv_topology import build_simple_uv_topology +from SkingToolkit.renderer import DifferentiableRenderer + + +def uv_structure_loss(logits, labels, renderer, views): + """Object-labelled UV graph edges, including physical cube seams.""" + p = logits.float().softmax(1)[:,1:] + g = F.one_hot(labels,4).permute(0,3,1,2)[:,1:].float() + groups = logits.shape[0] // len(views) + pred = p.new_zeros(groups,3,4096); truth = torch.zeros_like(pred) + count = p.new_zeros(groups,1,4096) + for vi,view in enumerate(views): + static = build_static_surface_routing(renderer,view,logits.device) + valid = static['masks'][1] & (static['part'][1]==0) + index = static['flat_uv'][1][valid] + pred.index_add_(2,index,p[vi::len(views),:,valid]) + truth.index_add_(2,index,g[vi::len(views),:,valid]) + count.index_add_(2,index,torch.ones(groups,1,index.numel(),device=p.device)) + pred, truth = pred/count.clamp_min(1), truth/count.clamp_min(1) + topology = build_simple_uv_topology() + edges = topology.outer_edge_index.to(p.device) + head = topology.part.flatten().to(p.device)==0 + a,b = edges[:,head[edges[0]] & head[edges[1]]] + visible = (count[:,:,a]>0) & (count[:,:,b]>0) + difference = (pred[:,:,a]-pred[:,:,b])-(truth[:,:,a]-truth[:,:,b]) + edge_loss = (difference.square()*visible).sum()/visible.expand_as(difference).sum().clamp_min(1) + support = count>0 + uv_loss = ((pred-truth).square()*support).sum()/support.expand_as(pred).sum().clamp_min(1) + faces=topology.face.flatten().to(p.device) + seam = faces[a] != faces[b] + seam_valid=visible[:,:,seam] + seam_loss=(difference[:,:,seam].square()*seam_valid).sum()/seam_valid.expand_as(difference[:,:,seam]).sum().clamp_min(1) + return uv_loss + edge_loss + seam_loss + + +@torch.no_grad() +def evaluate(model, renderer, dataset, config, output_dir, step, clean_paths, max_batches=None): + model.eval() + views = config['routing']['views'] + totals = {k:0 for k in ('tp','fp','fn','complete_objects','objects','uv_tp','uv_fn','head_rgb_abs','head_rgb_n','band_rgb_abs','band_rgb_n')} + per_class = {str(k):{'tp':0,'fp':0,'fn':0} for k in (1,2,3)} + visible_uv = torch.zeros(4096,dtype=torch.bool,device='cuda') + for view in views: + static = build_static_surface_routing(renderer,view,'cuda') + visible_uv[static['flat_uv'][1][static['masks'][1]]] = True + loader = DataLoader(dataset,batch_size=4,num_workers=0) + stress_rng=torch.Generator(device='cuda').manual_seed(10199) + for bi,batch in enumerate(loader): + if max_batches is not None and bi>=max_batches: break + uv,objects = batch['uv'].cuda(),batch['objects'].cuda() + images,labels,fg,components = render_accessories(uv,objects,renderer,views,batch['components'].cuda()) + if bi%2: + images=appearance_augment(images,fg[:,None],strength=1.25,generator=stress_rng) + result = run_pipeline(model,renderer,images,config,complete=True,foreground_probability=images[:,3] if config.get('synthetic_alpha_foreground') else None) + logits = result['outputs']['accessory_logits'] + cls = logits.argmax(1); confidence=logits.softmax(1).amax(1) + pred = result['details']['routing']['accessory_supported'] + gt = labels>0 + for k in (1,2,3): + pp,gg = pred&(cls==k),labels==k + for name,value in [('tp',pp&gg),('fp',pp&~gg),('fn',~pp&gg)]: per_class[str(k)][name]+=int(value.sum()) + for n in range(labels.shape[0]): + if gg[n].any(): + totals['objects']+=1 + totals['complete_objects']+=int((pp[n]&gg[n]).sum()/gg[n].sum() >= .95) + totals['tp']+=int((pred>).sum());totals['fp']+=int((pred&~gt).sum());totals['fn']+=int((~pred>).sum()) + uv_gt=(objects>0)&visible_uv.reshape(64,64) + known=result['conditioning'][:,9]>.5 + totals['uv_tp']+=int((known&uv_gt).sum());totals['uv_fn']+=int((~known&uv_gt).sum()) + headmask=torch.zeros_like(fg); y0,y1,x0,x1=head_bounds(*fg.shape[-2:]);headmask[:,y0:y1,x0:x1]=fg[:,y0:y1,x0:x1] + totals['head_rgb_abs']+=float((abs(result['render'][:,:3]-images[:,:3])*headmask[:,None]).sum()) + totals['head_rgb_n']+=int(headmask.sum())*3 + band=(components==2)|(components==4) + totals['band_rgb_abs']+=float((abs(result['render'][:,:3]-images[:,:3])*band[:,None]).sum()) + totals['band_rgb_n']+=int(band.sum())*3 + if bi==0: + save_image(images[:8,:3].cpu(),output_dir/f'validation_{step}_input.png',nrow=4) + save_image(result['render'][:8,:3].cpu(),output_dir/f'validation_{step}_render.png',nrow=4) + palette=images.new_tensor([[.5,.5,.5],[.1,.9,.5],[.95,.6,.1],[.7,.3,.9]]) + save_image(palette[cls[:8]].permute(0,3,1,2).cpu(),output_dir/f'validation_{step}_objects.png',nrow=4) + # Stored UV layers are a conservative regression guard, not semantic class GT. + from SkingToolkit.dense_uv_parser.skin_dataset import load_skin + inner_pixels=added_wrong=0 + for start in range(0,len(clean_paths),4): + uv=torch.stack([load_skin(path) for path in clean_paths[start:start+4]]).cuda() + images=torch.stack([renderer.forward_view(uv,v) for v in views],1).flatten(0,1) + result=run_pipeline(model,renderer,images,config,foreground_probability=images[:,3] if config.get('synthetic_alpha_foreground') else None) + base_outputs={k:v for k,v in result['outputs'].items() if k not in ('accessory_logits','hat_component_logits')} + baseline=run_pipeline(model,renderer,images,config,outputs=base_outputs,foreground_probability=images[:,3] if config.get('synthetic_alpha_foreground') else None) + topology=build_simple_uv_topology() + head_outer=(topology.valid&(topology.layer==1)&(topology.part==0)).cuda() + negative=(uv[:,3]<.5)&head_outer&visible_uv.reshape(64,64) + new=(result['conditioning'][:,9]>.5)&~(baseline['conditioning'][:,9]>.5) + inner_pixels+=int(negative.sum());added_wrong+=int((new&negative).sum()) + metrics={**totals,'classes':per_class,'validation_inputs':'Alternating canonical and stressed images; renderer alpha for synthetic foreground when enabled; frozen learned foreground on real regression', + 'object_precision':totals['tp']/max(1,totals['tp']+totals['fp']), + 'object_recall':totals['tp']/max(1,totals['tp']+totals['fn']), + 'object_iou':totals['tp']/max(1,totals['tp']+totals['fp']+totals['fn']), + 'complete_object_rate':totals['complete_objects']/max(1,totals['objects']), + 'hard_uv_object_recall':totals['uv_tp']/max(1,totals['uv_tp']+totals['uv_fn']), + 'hat_band_rgb_mae':totals['band_rgb_abs']/max(1,totals['band_rgb_n']), + 'final_head_rgb_mae':totals['head_rgb_abs']/max(1,totals['head_rgb_n']), + 'clean_added_outer_false_positive_rate':added_wrong/max(1,inner_pixels), + 'clean_negative_uv_texels':inner_pixels,'clean_added_wrong_outer':added_wrong} + from SkingToolkit.dense_uv_parser.accessory_regression import evaluate_real_regression + metrics['real_development'] = evaluate_real_regression(model,renderer,config,output_dir,step) + return metrics + + +def main(): + root=Path(__file__).resolve().parents[1] + parser=argparse.ArgumentParser(description=__doc__) + parser.add_argument('--output-dir',type=Path,required=True) + parser.add_argument('--steps',type=int,default=20000) + parser.add_argument('--eval-every',type=int,default=1000) + parser.add_argument('--batch-size',type=int,default=8) + parser.add_argument('--val-samples',type=int,default=64) + parser.add_argument('--clean-samples',type=int,default=32) + parser.add_argument('--resume',type=Path) + parser.add_argument('--hat-components',action='store_true') + parser.add_argument('--evaluate-only',action='store_true') + parser.add_argument('--lr',type=float,default=2e-4) + parser.add_argument('--pipeline',type=Path) + parser.add_argument('--replay-weight',type=float,default=.25) + parser.add_argument('--require-approved-start',action='store_true') + opt=parser.parse_args() + if min(opt.steps,opt.eval_every,opt.batch_size,opt.val_samples,opt.clean_samples)<=0 or opt.lr<=0 or opt.replay_weight<0: + parser.error('Counts and learning rate must be positive; replay weight nonnegative') + opt.output_dir.mkdir(parents=True,exist_ok=False) + torch.manual_seed(101);torch.set_num_threads(4) + torch.backends.cudnn.benchmark=True + torch.set_float32_matmul_precision('high') + config=load_pipeline(opt.pipeline);views=config['routing']['views'] + origin=root/'dense_uv_parser/runs/dense_uv_parser_v61/best.pt' + ckpt=torch.load(origin,map_location='cpu',weights_only=False) + kwargs={k:v for k,v in ckpt['model_config'].items() if k in inspect.signature(DenseUVParserNet).parameters} + kwargs['predict_head_accessories']=True + kwargs['predict_hat_components']=opt.hat_components + model=DenseUVParserNet(**kwargs).cuda() + missing,extra=model.load_state_dict(ckpt['model'],strict=False) + if extra or any(not k.startswith('accessory_head.') for k in missing): raise RuntimeError((missing,extra)) + if opt.resume: + missing,extra=model.load_state_dict(torch.load(opt.resume,map_location='cuda',weights_only=False)['model'],strict=False) + if extra or any(not k.startswith('accessory_head.components.') for k in missing):raise ValueError((missing,extra)) + for name,value in ckpt['model'].items(): + if not torch.equal(model.state_dict()[name].cpu(),value): + raise ValueError('Warm start changed frozen v61 weights: '+name) + model.requires_grad_(False);model.accessory_head.requires_grad_(True) + attach_semantic_runtime(model,'siglip2','google/siglip2-base-patch16-224','cuda',local_files_only=True) + mappings_dir=str(root.parent/'github/differentiable_minecraft_renderer/mappings_256x512') + renderer=DifferentiableRenderer(mappings_dir).cuda() + paths=SkinUVDataset(root/'skins',max_samples=180000).skin_paths + split=json.loads((root/'dense_uv_parser/runs/dense_uv_parser_v100/config.json').read_text())['experiment'] + train_paths=[paths[i] for i in split['train_indices']] + val_paths=[paths[i] for i in split['val_indices']] + test_paths=[paths[i] for i in split['test_indices']] + assert not (set(train_paths)&(set(val_paths)|set(test_paths))) + assert not (set(val_paths)&set(test_paths)) + source_splits={name:[str(p) for p in items] for name,items in + [('train',train_paths),('validation',val_paths),('test',test_paths)]} + split_bytes=json.dumps(source_splits,sort_keys=True).encode() + (opt.output_dir/'source_splits.json').write_bytes(split_bytes) + train=AccessoryDataset(train_paths,32768,101000) + val=AccessoryDataset(val_paths,opt.val_samples,9101000) + test=AccessoryDataset(test_paths,opt.val_samples,19101000) + loader=DataLoader(train,batch_size=opt.batch_size,shuffle=True,num_workers=4,pin_memory=True) + optimizer=torch.optim.AdamW(model.accessory_head.parameters(),lr=opt.lr,weight_decay=1e-4) + manifest={'version':'v101','base_checkpoint':str(origin),'base_sha256':hashlib.sha256(origin.read_bytes()).hexdigest(), + 'git_commit':subprocess.check_output(['git','rev-parse','HEAD'],cwd=root,text=True).strip(), + 'options':{k:str(v) if isinstance(v,Path) else v for k,v in vars(opt).items()}, + 'source_splits_sha256':hashlib.sha256(split_bytes).hexdigest(), + 'train_source_count':len(train_paths),'val_source_count':len(val_paths),'test_source_count':len(test_paths), + 'warm_start_sha256':hashlib.sha256(opt.resume.read_bytes()).hexdigest() if opt.resume else None, + 'training':'Only accessory branch is trainable; v61 body and original route weights are frozen. Warm starts use a fresh optimizer.', + 'labels':'Authored procedural object identities. Optional stored-inner replay is layer evidence, NOT semantic truth; it can contradict painted-on hats/hair. Set replay-weight=0 for semantic experiments.', + 'real_sample_scope':'TWRLRRTHQP2UV368 is a development regression image, never an independent test.', + 'source_sha256':{f.name:hashlib.sha256(f.read_bytes()).hexdigest() for f in Path(__file__).parent.glob('*.py')}, + 'pipeline':config,'classes':['none','glasses','hat','outer_hair'],'hat_components':['none','inner_crown','inner_band','outer_crown','outer_band','outer_brim'] if opt.hat_components else None} + write_json(opt.output_dir/'config.json',manifest) + snapshot=opt.output_dir/'source';snapshot.mkdir() + for f in Path(__file__).parent.glob('*.py'): shutil.copy2(f,snapshot/f.name) + (snapshot/'v101_pipeline.json').write_text(json.dumps(config,indent=2)) + shutil.copy2(Path(__file__).with_name('v101_real_regression.json'),snapshot/'v101_real_regression.json') + shutil.copytree(Path(__file__).with_name('regression'),snapshot/'regression') + start=time.time();best=-float('inf');step=0;epoch=0 + def save(name,metrics): + args={**ckpt['args'],'siglip_local_files_only':True,'outer_uv_min_source_pixels':30,'mappings_dir':mappings_dir,'views':views} + config_model={**ckpt['model_config'],'predict_head_accessories':True,'predict_hat_components':opt.hat_components,'accessory_route_threshold':config['accessory_route_threshold']} + payload={'model':model.state_dict(),'model_config':config_model,'args':args,'epoch':epoch,'step':step, + 'metrics':{'val':metrics},'v101_manifest':manifest,'optimizer':optimizer.state_dict()} + temp=opt.output_dir/(name+'.tmp');torch.save(payload,temp);os.replace(temp,opt.output_dir/name) + def admission(metrics): + real=metrics['real_development'] + return (metrics['object_precision']>=.90 and metrics['object_recall']>=.75 + and metrics['complete_object_rate']>=.25 + and metrics['clean_added_outer_false_positive_rate']<=.03 + and real['glasses_mask_precision']>=.85 + and real['whole_glasses_outer_recall']>=max(.90,real['baseline_whole_glasses_outer_recall']+.15) + and (not opt.hat_components or (min(real['hat_brim_outer_recall'])>=.90 and min(real['hat_band_chroma_retention'])>=.95))) + try: + if opt.evaluate_only: + metrics=evaluate(model,renderer,val,config,opt.output_dir,'calibrated',val_paths[:opt.clean_samples]) + accepted=admission(metrics) + save('latest.pt',metrics) + if accepted:save('best.pt',metrics) + write_json(opt.output_dir/'last_evaluation.json',{'step':0,'metrics':metrics,'accepted':accepted}) + write_json(opt.output_dir/'status.json',{'state':'evaluated','accepted':accepted}) + print(json.dumps(metrics),flush=True) + return + if opt.require_approved_start: + initial=evaluate(model,renderer,val,config,opt.output_dir,'initial',val_paths[:opt.clean_samples]) + write_json(opt.output_dir/'initial_evaluation.json',initial) + if not admission(initial):raise RuntimeError('Warm start failed release validation') + best=.4*initial['object_iou']+.3*initial['complete_object_rate']+.3*initial['hard_uv_object_recall'] + save('best.pt',initial) + print('approved_start='+json.dumps(initial),flush=True) + while step=opt.steps: break + model.eval();model.accessory_head.train() + images,labels,fg,components=render_accessories(batch['uv'].cuda(non_blocking=True),batch['objects'].cuda(non_blocking=True),renderer,views,batch['components'].cuda(non_blocking=True)) + if step%2==0: images=appearance_augment(images,fg[:,None],strength=1.25) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + prediction=model.predict_accessories(images,fg,return_components=opt.hat_components) + logits,component_logits=prediction if opt.hat_components else (prediction,None) + valid=torch.zeros_like(fg);y0,y1,x0,x1=head_bounds(*fg.shape[-2:]);valid[:,y0:y1,x0:x1]=True + loss=accessory_loss(logits,labels,valid)+.25*uv_structure_loss(logits,labels,renderer,views) + if component_logits is not None: + ce=F.cross_entropy(component_logits.float(),components.masked_fill(~valid,-100),weight=images.new_tensor([1.,1.5,3.,1.5,3.,4.])) + cp=component_logits.float().softmax(1)*valid[:,None] + cg=F.one_hot(components,6).permute(0,3,1,2).float()*valid[:,None] + dice=1-(2*(cp*cg).sum((2,3))+1)/(cp.sum((2,3))+cg.sum((2,3))+1) + loss=loss+ce+.5*dice[:,1:].mean() + if step%2==0 and opt.replay_weight>0: + # Only verified INNER head pixels are negative layer evidence. + # Unannotated outer texels are ignored, not assigned object IDs. + replay_uv=torch.stack([load_skin(train_paths[i]) for i in torch.randint(len(train_paths),(2,)).tolist()]).cuda() + replay_images=[];replay_fg=[];replay_valid=[] + for view in views: + im,t=build_dense_parser_batch(replay_uv,renderer,view) + replay_images.append(im);replay_fg.append(t['foreground'][:,0]) + replay_valid.append((t['layer']==0)&(t['part']==0)&(t['foreground'][:,0]>.5)) + ri=torch.stack(replay_images,1).flatten(0,1) + rf=torch.stack(replay_fg,1).flatten(0,1) + rv=torch.stack(replay_valid,1).flatten(0,1) + rl=model.predict_accessories(ri,rf) + if rv.any(): + loss=loss+opt.replay_weight*(-F.log_softmax(rl.float(),1)[:,0][rv]).mean() + if not torch.isfinite(loss): raise RuntimeError('Non-finite loss') + loss.backward();grad=torch.nn.utils.clip_grad_norm_(model.accessory_head.parameters(),1.) + if not torch.isfinite(grad): raise RuntimeError('Non-finite gradient') + optimizer.step();step+=1 + lr=opt.lr*(.1+.9*(1+math.cos(math.pi*step/opt.steps))/2) + for group in optimizer.param_groups: group['lr']=lr + if step%25==0 or step==1: + status={'state':'training','pid':os.getpid(),'step':step,'total_steps':opt.steps,'loss':float(loss.detach()), + 'gradient_norm':float(grad),'lr':lr,'elapsed_seconds':time.time()-start} + write_json(opt.output_dir/'status.json',status);print(json.dumps(status),flush=True) + if step%opt.eval_every==0 or step==opt.steps: + metrics=evaluate(model,renderer,val,config,opt.output_dir,step,val_paths[:opt.clean_samples]) + score=.4*metrics['object_iou']+.3*metrics['complete_object_rate']+.3*metrics['hard_uv_object_recall'] + accepted=admission(metrics) + save('latest.pt',metrics) + promoted=accepted and score>best + if promoted: best=score;save('best.pt',metrics) + record={'step':step,'metrics':metrics,'score':score,'accepted':accepted,'promoted':promoted} + write_json(opt.output_dir/'last_evaluation.json',record) + with (opt.output_dir/'metrics.jsonl').open('a') as f:f.write(json.dumps(record)+'\n') + print('evaluation='+json.dumps(record),flush=True) + epoch+=1 + for name,value in ckpt['model'].items(): + if not torch.equal(model.state_dict()[name].cpu(),value):raise RuntimeError('Frozen v61 parameter changed: '+name) + selected=opt.output_dir/'best.pt' + if selected.exists(): + model.load_state_dict(torch.load(selected,map_location='cuda',weights_only=False)['model']) + heldout=evaluate(model,renderer,test,config,opt.output_dir,'heldout',test_paths[:opt.clean_samples]) + write_json(opt.output_dir/'heldout_test.json',heldout) + write_json(opt.output_dir/'status.json',{'state':'complete','step':step,'best_checkpoint_exists':selected.exists(), + 'elapsed_seconds':time.time()-start}) + except BaseException as error: + write_json(opt.output_dir/'status.json',{'state':'failed','step':step,'error':repr(error)}) + raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train_foreground.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_foreground.py new file mode 100644 index 0000000000000000000000000000000000000000..686cdc9b5466f3ac9ff8fbb0f7b404ac17d48ce3 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_foreground.py @@ -0,0 +1,198 @@ +"""Supervised Minecraft foreground decoder adaptation with independent identity splits.""" +import argparse,hashlib,json,math,os,shutil,subprocess,time +from pathlib import Path +from types import SimpleNamespace +import numpy as np +from PIL import Image,ImageFilter +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from torchvision.utils import save_image +from SkingToolkit.dense_uv_parser.matting_data import MattingDataset +from SkingToolkit.dense_uv_parser.foreground_model import load_foreground_model,foreground_logits,sha + + +def write_json(path,value): + p=Path(path);tmp=p.with_suffix(p.suffix+'.tmp');tmp.write_text(json.dumps(value,indent=2,allow_nan=False));os.replace(tmp,p) + + +def boundary(mask,width=2): + dil=F.max_pool2d(mask.float(),width*2+1,1,width) + ero=-F.max_pool2d(-mask.float(),width*2+1,1,width) + return (dil-ero)>.01 + + +def alpha_loss(logits,target): + p=logits.float().sigmoid();edge=boundary(target) + weight=1+4*edge.float() + bce=(F.binary_cross_entropy_with_logits(logits.float(),target,reduction='none')*weight).sum()/weight.sum() + dice=1-(2*(p*target).sum((1,2,3))+1)/(p.sum((1,2,3))+target.sum((1,2,3))+1) + edge_error=(abs(p-target)*edge).sum()/edge.sum().clamp_min(1) + grad=sum(abs(torch.diff(p,dim=d)-torch.diff(target,dim=d)).mean() for d in (2,3)) + # Mine the worst background pixels, including wholly empty crops. + # Boundary-only improvement must not hide false foreground islands. + bg_loss=F.softplus(logits.float())*(target<.01).float() + hard_bg=bg_loss.flatten(1).topk(max(1,bg_loss[0].numel()//100),dim=1).values.mean() + return bce+.5*dice.mean()+2*edge_error+.2*grad+.5*hard_bg + + +@torch.no_grad() +def evaluate(model,manifest,split,count,size,out,tag,real=True): + model.eval();dataset=MattingDataset(manifest,split,size=size,limit=count) + loader=DataLoader(dataset,batch_size=2,num_workers=4) + totals={k:0 for k in ('tp','fp','fn','boundary_abs','boundary_n','abs','n','edge_tp_p','edge_tp_g','edge_p','edge_g')} + by_background={} + for bi,batch in enumerate(loader): + x,a=batch['image'].cuda(),batch['alpha'].cuda() + with torch.autocast('cuda',dtype=torch.bfloat16):p=foreground_logits(model,x).float().sigmoid() + pred=p>=.5;gt=a>=.5 + totals['tp']+=int((pred>).sum());totals['fp']+=int((pred&~gt).sum());totals['fn']+=int((~pred>).sum()) + edge=boundary(a) + totals['boundary_abs']+=float((abs(p-a)*edge).sum());totals['boundary_n']+=int(edge.sum()) + totals['abs']+=float(abs(p-a).sum());totals['n']+=p.numel() + ep=boundary(pred,1);eg=boundary(gt,1) + totals['edge_tp_p']+=int((ep&(F.max_pool2d(eg.float(),5,1,2)>.5)).sum()) + totals['edge_tp_g']+=int((eg&(F.max_pool2d(ep.float(),5,1,2)>.5)).sum()) + totals['edge_p']+=int(ep.sum());totals['edge_g']+=int(eg.sum()) + for i,k in enumerate(batch['background_kind'].tolist()): + row=by_background.setdefault(str(k),{'tp':0,'fp':0,'fn':0}) + row['tp']+=int((pred[i]>[i]).sum());row['fp']+=int((pred[i]&~gt[i]).sum());row['fn']+=int((~pred[i]>[i]).sum()) + if bi==0: + save_image(torch.cat([x,torch.where(pred,x,.5),a.expand(-1,3,-1,-1),p.expand(-1,3,-1,-1)],dim=0).cpu(),out/f'{tag}_validation.png',nrow=2) + precision=totals['edge_tp_p']/max(1,totals['edge_p']);recall=totals['edge_tp_g']/max(1,totals['edge_g']) + result={**totals,'images':len(dataset),'iou':totals['tp']/max(1,totals['tp']+totals['fp']+totals['fn']), + 'boundary_f1_tolerance_2px':2*precision*recall/max(1e-8,precision+recall), + 'boundary_alpha_mae':totals['boundary_abs']/max(1,totals['boundary_n']), + 'alpha_mae':totals['abs']/max(1,totals['n']), + 'background_iou':{k:v['tp']/max(1,sum(v.values())) for k,v in by_background.items()}} + if real:result['real_development']=evaluate_real(model,out,tag) + return result + + +@torch.no_grad() +def evaluate_real(model,out,tag): + root=Path(__file__).parent + cases=json.loads((root/'v101_cases.json').read_text()) + targets=[x for x in cases if any(k in x for k in ('PU418PLGER','TWRLRRTH','img28'))] + metrics={'role':'development checks, no gradient training on these images'} + dest=out/f'real_{tag}';dest.mkdir(exist_ok=True) + for source in targets: + path=Path(source);rgb=np.array(Image.open(path).convert('RGB')) + image=torch.from_numpy(rgb.copy()).permute(2,0,1).float().cuda()/255 + views=torch.stack(image.chunk(2,dim=2)) + resized=F.interpolate(views,(1024,1024),mode='bilinear',align_corners=False) + with torch.autocast('cuda',dtype=torch.bfloat16):prob=foreground_logits(model,resized).float().sigmoid() + native=F.interpolate(prob,(512,256),mode='nearest-exact')[:,0] + original=F.interpolate(prob,views.shape[-2:],mode='bilinear',align_corners=False)[:,0] + array=torch.cat(list(original),dim=1).cpu().numpy() + Image.fromarray((array*255).astype('uint8')).save(dest/(path.stem+'_probability.png')) + save_image(torch.where(native[:,None]>=.5,F.interpolate(views,(512,256),mode='nearest-exact'),.5).cpu(),dest/(path.stem+'_cutout.png'),nrow=2) + if 'PU418' in path.name: + mask_path=root/'runs/v101_hat_audit_20260906/hat_brim_development_mask.png' + mask=torch.from_numpy(np.array(Image.open(mask_path).filter(ImageFilter.MinFilter(3)))>0).cuda() + metrics['hat_brim_foreground_recall']=[float((v[mask]>=.5).float().mean()) for v in native] + if 'TWRLR' in path.name: + mask_path=root/'regression/TWRLRRTHQP2UV368_glasses.png' + mask=torch.from_numpy(np.array(Image.open(mask_path).filter(ImageFilter.MinFilter(5)))>0).cuda() + metrics['glasses_foreground_recall']=float((native[0][mask]>=.5).float().mean()) + return metrics + + +def main(): + root=Path(__file__).resolve().parents[1];p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--dataset',type=Path,required=True);p.add_argument('--output-dir',type=Path,required=True) + p.add_argument('--model-dir',type=Path,required=True);p.add_argument('--resume',type=Path) + p.add_argument('--steps',type=int,default=800);p.add_argument('--eval-every',type=int,default=200) + p.add_argument('--batch-size',type=int,default=4);p.add_argument('--train-size',type=int,default=512) + p.add_argument('--eval-size',type=int,default=768);p.add_argument('--val-samples',type=int,default=96) + p.add_argument('--lr',type=float,default=0.000005) + opt=p.parse_args() + if min(opt.steps,opt.eval_every,opt.batch_size,opt.val_samples)<=0:raise ValueError('Counts must be positive') + opt.output_dir.mkdir(parents=True,exist_ok=False) + os.environ['HF_HUB_OFFLINE']='1';os.environ['HF_HUB_DISABLE_PROGRESS_BARS']='1' + torch.manual_seed(106);torch.set_num_threads(4);torch.backends.cudnn.benchmark=True + model,_=load_foreground_model(opt.model_dir,opt.resume) + for name,param in model.named_parameters():param.requires_grad_(not name.startswith('bb.')) + trainable=[v for v in model.parameters() if v.requires_grad] + optimizer=torch.optim.AdamW(trainable,lr=opt.lr,weight_decay=1e-4) + def backbone_hash(): + h=hashlib.sha256() + for name,value in model.state_dict().items(): + if name.startswith('bb.'):h.update(name.encode());h.update(value.cpu().numpy().tobytes()) + return h.hexdigest() + frozen_hash=backbone_hash() + info={'version':'minecraft_foreground_v1','architecture':'BiRefNet pretrained frozen backbone; supervised decoder adaptation', + 'base_model_dir':str(opt.model_dir.resolve()),'base_sha256':sha(opt.model_dir/'model.safetensors'), + 'backbone_tensor_sha256':frozen_hash,'dataset_manifest_sha256':sha(opt.dataset), + 'dataset_manifest':str(opt.dataset.resolve()),'trainable_parameters':sum(x.numel() for x in trainable), + 'options':{k:str(v) if isinstance(v,Path) else v for k,v in vars(opt).items()}, + 'git_commit':subprocess.check_output(['git','rev-parse','HEAD'],cwd=root,text=True).strip(), + 'sources':{f.name:sha(f) for f in Path(__file__).parent.glob('*.py')}, + 'supervision':'Renderer alpha from independent normalized skin identities; randomized backgrounds and premultiplied affine resampling. No real regression images used for gradients.', + 'inference_resolution':1024,'output':'Supervised alpha prediction; retain separate reliable RGB source mask'} + write_json(opt.output_dir/'config.json',info) + snapshot=opt.output_dir/'source';snapshot.mkdir() + for f in Path(__file__).parent.glob('*.py'):shutil.copy2(f,snapshot/f.name) + step=0;best=float('inf');start=time.time() + def save(name,metrics): + decoder={k:v.detach().cpu() for k,v in model.state_dict().items() if not k.startswith('bb.')} + temp=opt.output_dir/(name+'.tmp') + torch.save({'decoder':decoder,'optimizer':optimizer.state_dict(),'step':step,'metrics':metrics,'manifest':info},temp) + os.replace(temp,opt.output_dir/name) + def accepted(m,baseline): + r=m['real_development'] + return (m['iou']>=max(.985,baseline['iou']-.002) and m['boundary_f1_tolerance_2px']>=baseline['boundary_f1_tolerance_2px']-.003 + and m['boundary_alpha_mae']<=baseline['boundary_alpha_mae']*.98 + and min(r['hat_brim_foreground_recall']+[r['glasses_foreground_recall']])>=.995) + try: + write_json(opt.output_dir/'status.json',{'state':'baseline_validation','pid':os.getpid(),'step':0}) + baseline=evaluate(model,opt.dataset,'validation',opt.val_samples,opt.eval_size,opt.output_dir,'baseline') + write_json(opt.output_dir/'baseline.json',baseline);print('baseline='+json.dumps(baseline),flush=True) + dataset=MattingDataset(opt.dataset,'train',size=opt.train_size) + epoch=0 + while step=opt.steps:break + model.eval() # preserves pretrained BN statistics and uses the inference decoder path; gradients remain enabled + x,a=batch['image'].cuda(non_blocking=True),batch['alpha'].cuda(non_blocking=True) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16):loss=alpha_loss(foreground_logits(model,x),a) + if not torch.isfinite(loss):raise RuntimeError('Non-finite foreground loss') + loss.backward();gn=torch.nn.utils.clip_grad_norm_(trainable,1.) + if not torch.isfinite(gn):raise RuntimeError('Non-finite foreground gradient') + optimizer.step();step+=1 + rate=opt.lr*(.1+.9*(1+math.cos(math.pi*step/opt.steps))/2) + for group in optimizer.param_groups:group['lr']=rate + if step==1 or step%10==0: + state={'state':'training','pid':os.getpid(),'step':step,'total_steps':opt.steps,'loss':float(loss),'gradient_norm':float(gn),'elapsed_seconds':time.time()-start} + write_json(opt.output_dir/'status.json',state);print(json.dumps(state),flush=True) + if step%opt.eval_every==0 or step==opt.steps: + metrics=evaluate(model,opt.dataset,'validation',opt.val_samples,opt.eval_size,opt.output_dir,str(step)) + valid=accepted(metrics,baseline);score=metrics['boundary_alpha_mae']+5*(1-metrics['iou']) + save('latest.pt',metrics);promoted=valid and score.5,return_presence=True);gt=(labels==4).flatten(1).any(1);pred=prediction.sigmoid()>=.9 + tp+=int((pred>).sum());fp+=int((pred&~gt).sum());tn+=int((~pred&~gt).sum());fn+=int((~pred>).sum()) + return {'tp':tp,'fp':fp,'tn':tn,'fn':fn,'precision':tp/max(1,tp+fp),'recall':tp/max(1,tp+fn),'negative_fpr':fp/max(1,fp+tn)} + @torch.no_grad() + def real(): + model.eval();stats={} + for folder in sorted((root/'output_history/v101_retrained_20260906').iterdir()): + m=json.loads((folder/'manifest.json').read_text());d=torch.load(folder/'diagnostics.pt',weights_only=False) + images=d['images'].cuda();fg=cached_real_foreground(images,m['input'],pipeline) + _,prediction=model.predict_ownership(images,fg>=.5,return_presence=True) + stats[Path(m['input']).stem]=prediction.sigmoid().tolist() + return stats + start=time.time();step=0 + try: + for batch in DataLoader(train,batch_size=8,shuffle=True,num_workers=4,pin_memory=True): + if step>=opt.steps:break + model.eval();model.ownership_head.presence.train() + with torch.no_grad():images,labels=render_ownership(batch['uv'].cuda(),batch['ownership'].cuda(),renderer,['front_left','back_left']);fg=images[:,3]>.5 + if step%2==0:images=appearance_augment(images,fg[:,None],strength=1.4) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + _,prediction=model.predict_ownership(images,fg,return_presence=True);truth=(labels==4).flatten(1).any(1).float();loss=F.binary_cross_entropy_with_logits(prediction.float(),truth,weight=torch.where(truth>.5,1.,2.)) + if not torch.isfinite(loss):raise RuntimeError('Non-finite gate loss') + loss.backward();optimizer.step();step+=1 + if step%50==0: + status={'state':'training','step':step,'loss':float(loss.detach()),'elapsed_seconds':time.time()-start};write_json(out/'status.json',status);print(json.dumps(status),flush=True) + if step%300==0 or step==opt.steps: + metrics={'step':step,'validation':evaluate(val),'real':real()};write_json(out/f'evaluation_{step}.json',metrics) + payload={**{k:v for k,v in ckpt.items() if k!='model'},'model':model.state_dict(),'model_config':{**ckpt['model_config'],'predict_headphone_presence':True},'inference_pipeline':pipeline,'presence_manifest':manifest,'presence_metrics':metrics} + torch.save(payload,out/f'step_{step}.pt');print(json.dumps(metrics),flush=True) + for name,value in ckpt['model'].items(): + if not torch.equal(value,model.state_dict()[name].cpu()):raise ValueError('Frozen tensor changed: '+name) + write_json(out/'heldout.json',evaluate(test));write_json(out/'status.json',{'state':'complete','step':step,'elapsed_seconds':time.time()-start}) + except BaseException as e:write_json(out/'status.json',{'state':'failed','step':step,'error':repr(e)});raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear.py new file mode 100644 index 0000000000000000000000000000000000000000..de891110f4333feecdb3cdae0bafc84f2140481c --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear.py @@ -0,0 +1,114 @@ +"""Adapt an independent headwear decoder, preserving the released parser.""" +import argparse, hashlib, inspect, json, math, shutil, time +from pathlib import Path +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet +from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime +from SkingToolkit.dense_uv_parser.headwear_data import HeadwearDataset, render_headwear +from SkingToolkit.dense_uv_parser.headwear import HEADWEAR_CLASSES +from SkingToolkit.dense_uv_parser.accessories import head_bounds +from SkingToolkit.dense_uv_parser.semantic_generalization import appearance_augment, write_json +from SkingToolkit.renderer import DifferentiableRenderer + + +def headwear_loss(logits, truth): + probability = logits.float().softmax(1) + ce = F.cross_entropy(logits.float(), truth, weight=logits.new_tensor([1,2,4,2,4,3,4]).float()) + onehot = F.one_hot(truth,7).permute(0,3,1,2).float() + dice = 1-(2*(probability*onehot).sum((2,3))+1)/(probability.sum((2,3))+onehot.sum((2,3))+1) + # One authored component has one layer across the two views. Penalize the + # opposite layer on every member, including low confidence boundaries. + consistency = logits.sum()*0 + for ids in ((1,3),(2,4)): + for start in range(0,truth.shape[0],2): + mask=(truth[start:start+2]==ids[0])|(truth[start:start+2]==ids[1]) + if not mask.any():continue + correct=int((truth[start:start+2][mask]==ids[1]).all()) + pair=probability[start:start+2][:,ids] + logp=(pair[:,correct]/pair.sum(1).clamp_min(1e-6)).clamp_min(1e-6).log() + consistency += -logp[mask].mean()/max(1,truth.shape[0]) + return ce+dice[:,1:].mean()+.5*consistency + + +@torch.no_grad() +def evaluate(model,renderer,dataset): + confusion=torch.zeros(7,7,device='cuda',dtype=torch.long);model.eval() + for batch in DataLoader(dataset,batch_size=8): + images,truth=render_headwear(batch['uv'].cuda(),batch['components'].cuda(),renderer,['front_left','back_left']) + logits=model.predict_headwear_components(images,images[:,3]>.5) + y0,y1,x0,x1=head_bounds(*truth.shape[-2:]) + truth=truth[:,y0:y1,x0:x1];pred=logits.argmax(1)[:,y0:y1,x0:x1] + confusion+=torch.bincount((truth*7+pred).flatten(),minlength=49).reshape(7,7) + tp=confusion.diag() + return {'confusion':confusion.cpu().tolist(),'precision':(tp/confusion.sum(0).clamp_min(1)).cpu().tolist(),'recall':(tp/confusion.sum(1).clamp_min(1)).cpu().tolist()} + + +def main(): + p=argparse.ArgumentParser(description=__doc__) + p.add_argument('--output-dir',type=Path,required=True);p.add_argument('--steps',type=int,default=1800) + p.add_argument('--eval-every',type=int,default=600);p.add_argument('--batch-size',type=int,default=8) + p.add_argument('--lr',type=float,default=1e-4);p.add_argument('--resume',type=Path) + opt=p.parse_args();out=opt.output_dir;out.mkdir(parents=True,exist_ok=False) + root=Path(__file__).parent;origin=root/'runs/v101_semantic_release_gated_20260906/parser.pt' + torch.manual_seed(910107);torch.set_num_threads(4);torch.backends.cudnn.benchmark=True;torch.set_float32_matmul_precision('high') + ckpt=torch.load(origin,map_location='cpu',weights_only=False) + kwargs={k:v for k,v in ckpt['model_config'].items() if k in inspect.signature(DenseUVParserNet).parameters} + kwargs['predict_headwear']=True;model=DenseUVParserNet(**kwargs).cuda() + missing,extra=model.load_state_dict(ckpt['model'],strict=False) + if extra or any(not x.startswith('headwear_head.') for x in missing):raise ValueError((missing,extra)) + warm={k.removeprefix('accessory_head.'):v for k,v in ckpt['model'].items() if k.startswith('accessory_head.') and not k.startswith(('accessory_head.components.','accessory_head.classifier.'))} + model.headwear_head.load_state_dict(warm,strict=False) + with torch.no_grad(): + model.headwear_head.classifier.weight[:6].copy_(ckpt['model']['accessory_head.components.weight']) + model.headwear_head.classifier.bias[:6].copy_(ckpt['model']['accessory_head.components.bias']) + model.headwear_head.classifier.weight[6].zero_();model.headwear_head.classifier.bias[6]=-4 + if opt.resume:model.load_state_dict(torch.load(opt.resume,map_location='cuda',weights_only=False)['model']) + model.requires_grad_(False);model.headwear_head.requires_grad_(True) + attach_semantic_runtime(model,'siglip2','google/siglip2-base-patch16-224','cuda',local_files_only=True) + renderer=DifferentiableRenderer(ckpt['args']['mappings_dir']).cuda() + split=json.loads((root/'runs/dense_uv_parser_v101_retrain_20260906/source_splits.json').read_text()) + train=HeadwearDataset(split['train'],32768,7101000);val=HeadwearDataset(split['validation'],64,69101000) + assert not set(split['train'])&(set(split['validation'])|set(split['test'])) + manifest={'classes':HEADWEAR_CLASSES,'base_checkpoint':str(origin),'base_sha256':hashlib.sha256(origin.read_bytes()).hexdigest(), + 'options':{k:str(v) if isinstance(v,Path) else v for k,v in vars(opt).items()}, + 'training':'Authored whole components and visibility labels; all released parameters frozen; real examples are development only.', + 'source_sha256':{f.name:hashlib.sha256(f.read_bytes()).hexdigest() for f in root.glob('*.py')}} + write_json(out/'config.json',manifest);snapshot=out/'source';snapshot.mkdir() + for f in root.glob('*.py'):shutil.copy2(f,snapshot/f.name) + loader=DataLoader(train,batch_size=opt.batch_size,shuffle=True,num_workers=4,pin_memory=True) + optimizer=torch.optim.AdamW(model.headwear_head.parameters(),lr=opt.lr,weight_decay=1e-4) + start=time.time();step=0;epoch=0 + try: + while step=opt.steps:break + model.eval();model.headwear_head.train() + with torch.no_grad():images,truth=render_headwear(batch['uv'].cuda(),batch['components'].cuda(),renderer,['front_left','back_left']);fg=images[:,3]>.5 + if step%2==0:images=appearance_augment(images,fg[:,None],strength=1.4) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + logits=model.predict_headwear_components(images,fg);y0,y1,x0,x1=head_bounds(*truth.shape[-2:]) + loss=headwear_loss(logits[:,:,y0:y1,x0:x1],truth[:,y0:y1,x0:x1]) + if not torch.isfinite(loss):raise RuntimeError('Non-finite loss') + loss.backward();grad=torch.nn.utils.clip_grad_norm_(model.headwear_head.parameters(),1.) + if not torch.isfinite(grad):raise RuntimeError('Non-finite gradient') + optimizer.step();step+=1 + for group in optimizer.param_groups:group['lr']=opt.lr*(.1+.9*(1+math.cos(math.pi*step/opt.steps))/2) + if step%25==0 or step==1: + status={'state':'training','step':step,'total_steps':opt.steps,'loss':float(loss.detach()),'elapsed_seconds':time.time()-start};write_json(out/'status.json',status);print(json.dumps(status),flush=True) + if step%opt.eval_every==0 or step==opt.steps: + metrics=evaluate(model,renderer,val);write_json(out/f'evaluation_{step}.json',metrics) + for name,value in ckpt['model'].items(): + if not torch.equal(model.state_dict()[name].cpu(),value):raise RuntimeError('Released tensor changed: '+name) + payload={**{k:v for k,v in ckpt.items() if k not in ('model','optimizer')},'model':model.state_dict(),'model_config':{**ckpt['model_config'],'predict_headwear':True},'step':step,'headwear_manifest':manifest,'headwear_metrics':metrics} + torch.save(payload,out/f'step_{step}.pt');print('evaluation='+json.dumps(metrics),flush=True) + epoch+=1 + write_json(out/'status.json',{'state':'complete','step':step,'released_tensors_unchanged':True,'elapsed_seconds':time.time()-start}) + except BaseException as error: + write_json(out/'status.json',{'state':'failed','step':step,'error':repr(error)});raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear_presence.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear_presence.py new file mode 100644 index 0000000000000000000000000000000000000000..80a38acf5ed88fd1a80a8f51eeba244e70c08c47 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_headwear_presence.py @@ -0,0 +1,71 @@ +"""Train whole-headwear presence on authored positives and rich hair negatives.""" +from pathlib import Path +import argparse,inspect,json,hashlib,time,os +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet +from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime +from SkingToolkit.dense_uv_parser.headwear_data import HeadwearPresenceDataset,render_headwear +from SkingToolkit.dense_uv_parser.semantic_generalization import appearance_augment,write_json +from SkingToolkit.dense_uv_parser.accessory_pipeline import cached_real_foreground +from SkingToolkit.dense_uv_parser.accessories import head_bounds +from SkingToolkit.renderer import DifferentiableRenderer + + +def main(): + p=argparse.ArgumentParser();p.add_argument('--output-dir',type=Path,required=True);p.add_argument('--steps',type=int,default=1200);opt=p.parse_args();out=opt.output_dir;out.mkdir(parents=True,exist_ok=False) + root=Path(__file__).parent;torch.set_num_threads(4);torch.manual_seed(1010602) + source=root/'runs/v101_headwear_main_20260906/step_3600.pt';ckpt=torch.load(source,map_location='cpu',weights_only=False) + cfg={k:v for k,v in ckpt['model_config'].items() if k in inspect.signature(DenseUVParserNet).parameters};cfg['predict_headwear_presence']=True + model=DenseUVParserNet(**cfg).cuda();missing,extra=model.load_state_dict(ckpt['model'],strict=False) + if extra or any(not x.startswith('headwear_head.presence.') for x in missing):raise ValueError((missing,extra)) + model.requires_grad_(False);model.headwear_head.presence.requires_grad_(True) + attach_semantic_runtime(model,'siglip2','google/siglip2-base-patch16-224','cuda',local_files_only=True) + renderer=DifferentiableRenderer(ckpt['args']['mappings_dir']).cuda();pipeline=json.loads((root/'runs/v101_semantic_release_gated_20260906/pipeline.json').read_text()) + split=json.loads((root/'runs/dense_uv_parser_v101_retrain_20260906/source_splits.json').read_text()) + train=HeadwearPresenceDataset(split['train'],32768,95101000);val=HeadwearPresenceDataset(split['validation'],256,105101000);test=HeadwearPresenceDataset(split['test'][192:256],128,115101000) + optimizer=torch.optim.AdamW(model.headwear_head.presence.parameters(),lr=.001,weight_decay=.001) + manifest={'parent':str(source),'parent_sha256':hashlib.sha256(source.read_bytes()).hexdigest(),'method':'Only whole-object gate trained on pooled component maps and frozen SigLIP features. Rich authored hair, hats and headphones are negatives for royal crowns. All previous weights remain frozen.','steps':opt.steps,'pipeline':pipeline,'source_sha256':{f.name:hashlib.sha256(f.read_bytes()).hexdigest() for f in root.glob('*.py')}};write_json(out/'config.json',manifest) + @torch.no_grad() + def evaluate(dataset): + model.eval();tp=fp=tn=fn=0 + for batch in DataLoader(dataset,batch_size=16): + images,labels=render_headwear(batch['uv'].cuda(),batch['components'].cuda(),renderer,['front_left','back_left']) + _,prediction=model.predict_headwear_components(images,images[:,3]>.5,return_presence=True);gt=torch.stack([((labels>=1)&(labels<=5)).flatten(1).any(1),(labels==6).flatten(1).any(1)],1);pred=prediction.sigmoid()>=.95 + tp+=int((pred>).sum());fp+=int((pred&~gt).sum());tn+=int((~pred&~gt).sum());fn+=int((~pred>).sum()) + return {'tp':tp,'fp':fp,'tn':tn,'fn':fn,'precision':tp/max(1,tp+fp),'recall':tp/max(1,tp+fn),'negative_fpr':fp/max(1,fp+tn)} + @torch.no_grad() + def real(): + model.eval();stats={} + for folder in sorted((root/'output_history/v101_retrained_20260906').iterdir()): + m=json.loads((folder/'manifest.json').read_text());d=torch.load(folder/'diagnostics.pt',weights_only=False) + images=d['images'].cuda();fg=cached_real_foreground(images,m['input'],pipeline) + _,prediction=model.predict_headwear_components(images,fg>=.5,return_presence=True) + stats[Path(m['input']).stem]=prediction.sigmoid().tolist() + return stats + start=time.time();step=0 + try: + for batch in DataLoader(train,batch_size=8,shuffle=True,num_workers=4,pin_memory=True): + if step>=opt.steps:break + model.eval();model.headwear_head.presence.train() + with torch.no_grad():images,labels=render_headwear(batch['uv'].cuda(),batch['components'].cuda(),renderer,['front_left','back_left']);fg=images[:,3]>.5 + if step%2==0:images=appearance_augment(images,fg[:,None],strength=1.4) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + _,prediction=model.predict_headwear_components(images,fg,return_presence=True);truth=torch.stack([((labels>=1)&(labels<=5)).flatten(1).any(1),(labels==6).flatten(1).any(1)],1).float();loss=F.binary_cross_entropy_with_logits(prediction.float(),truth,weight=torch.where(truth>.5,1.,2.)) + if not torch.isfinite(loss):raise RuntimeError('Non-finite gate loss') + loss.backward();optimizer.step();step+=1 + if step%50==0: + status={'state':'training','step':step,'loss':float(loss.detach()),'elapsed_seconds':time.time()-start};write_json(out/'status.json',status);print(json.dumps(status),flush=True) + if step%300==0 or step==opt.steps: + metrics={'step':step,'validation':evaluate(val),'real':real()};write_json(out/f'evaluation_{step}.json',metrics) + payload={**{k:v for k,v in ckpt.items() if k!='model'},'model':model.state_dict(),'model_config':{**ckpt['model_config'],'predict_headwear_presence':True},'inference_pipeline':pipeline,'presence_manifest':manifest,'presence_metrics':metrics} + torch.save(payload,out/f'step_{step}.pt');print(json.dumps(metrics),flush=True) + for name,value in ckpt['model'].items(): + if not torch.equal(value,model.state_dict()[name].cpu()):raise ValueError('Frozen tensor changed: '+name) + write_json(out/'heldout.json',evaluate(test));write_json(out/'status.json',{'state':'complete','step':step,'elapsed_seconds':time.time()-start}) + except BaseException as e:write_json(out/'status.json',{'state':'failed','step':step,'error':repr(e)});raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/train_ownership.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_ownership.py new file mode 100644 index 0000000000000000000000000000000000000000..a75e6485131c3836db4f24281320b532892107be --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/train_ownership.py @@ -0,0 +1,122 @@ +"""Train head ownership independently, preserving every released v101 parameter.""" +import argparse,hashlib,inspect,json,os,shutil,time,math +from pathlib import Path +import torch +import torch.nn.functional as F +from torch.utils.data import DataLoader +from torchvision.utils import save_image +from SkingToolkit.dense_uv_parser.model import DenseUVParserNet +from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime +from SkingToolkit.dense_uv_parser.ownership_data import OwnershipDataset,render_ownership +from SkingToolkit.dense_uv_parser.ownership import OWNERSHIP_CLASSES +from SkingToolkit.dense_uv_parser.accessories import head_bounds +from SkingToolkit.dense_uv_parser.accessory_pipeline import run_pipeline,cached_real_foreground +from SkingToolkit.dense_uv_parser.semantic_generalization import appearance_augment,write_json +from SkingToolkit.dense_uv_parser.uv_layout import tensor_to_rgba_image +from SkingToolkit.renderer import DifferentiableRenderer + + +@torch.no_grad() +def evaluate(model,renderer,dataset): + confusion=torch.zeros(5,5,device='cuda',dtype=torch.long);model.eval() + for batch in DataLoader(dataset,batch_size=8): + images,labels=render_ownership(batch['uv'].cuda(),batch['ownership'].cuda(),renderer,['front_left','back_left']) + logits=model.predict_ownership(images,images[:,3]>.5);y0,y1,x0,x1=head_bounds(*labels.shape[-2:]) + truth=labels[:,y0:y1,x0:x1];pred=logits.argmax(1)[:,y0:y1,x0:x1] + confusion+=torch.bincount((truth*5+pred).flatten(),minlength=25).reshape(5,5) + tp=confusion.diag();precision=tp/confusion.sum(0).clamp_min(1);recall=tp/confusion.sum(1).clamp_min(1) + return {'confusion':confusion.cpu().tolist(),'precision':precision.cpu().tolist(),'recall':recall.cpu().tolist()} + + +@torch.no_grad() +def real_review(model,renderer,config,out,step): + cases=[('glasses','TWRLRRTHQP2UV368_edited_1f54edecd442'),('headphones','img27_610024a091fe'),('hat','PU418PLGERMFC487_edited_edada35d51d2')] + root=Path(__file__).parent + stats={} + for name,folder in cases: + src=root/'output_history/v101_retrained_20260906'/folder + manifest=json.loads((src/'manifest.json').read_text());d=torch.load(src/'diagnostics.pt',weights_only=False) + images=d['images'].cuda();fg=cached_real_foreground(images,manifest['input'],config) + outputs={k:v.cuda() if torch.is_tensor(v) else v for k,v in d['outputs'].items()} + outputs['head_ownership_logits']=model.predict_ownership(images,fg>=.5) + r=run_pipeline(model,renderer,images,config,complete=True,outputs=outputs,foreground_probability=fg) + p=out/f'real_{step}'/name;p.mkdir(parents=True,exist_ok=True) + tensor_to_rgba_image(r['uv'][0]).save(p/'uv.png');save_image(torch.cat(list(r['render'][:,:3]),2),p/'render.png') + palette=images.new_tensor([[.5,.5,.5],[1.,.65,.4],[.4,.3,.8],[.1,.9,.5],[1.,.2,.3]]) + cls=outputs['head_ownership_logits'].argmax(1) + save_image(torch.cat(list(palette[cls].permute(0,3,1,2)),2),p/'ownership.png') + route=r['details']['routing'] + torch.save({'uv':r['uv'].cpu(),'conditioning':r['conditioning'].cpu(),'outputs':{k:v.cpu() if torch.is_tensor(v) else v for k,v in outputs.items()},'routing':{k:v.cpu() if torch.is_tensor(v) else v for k,v in route.items()}},p/'diagnostics.pt') + # Development probes are not labels or gradients; colour is only a diagnostic. + uv=r['uv'][0];stats[name]={'inner_corrections':int(route['ownership_inner_supported'].sum()),'headphone_pixels':int(route['ownership_outer_supported'].sum())} + if name=='glasses':stats[name]['forehead_outer_alpha']=float(uv[3,10,41]) + if name=='headphones': + green=(uv[1]-torch.maximum(uv[0],uv[2])).clamp_min(0) + mask=torch.zeros_like(green,dtype=torch.bool);mask[:8,8:16]=True;mask[8:16,:8]=True;mask[8:16,16:32]=True + stats[name]['inner_nonface_green_sum']=float(green[mask].sum()) + if name=='hat':stats[name]['brim_ring']=int(sum((uv[3,11,x:x+8]>.5).sum() for x in [32,40,48,56])) + return stats + + +def main(): + parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output-dir',type=Path,required=True);parser.add_argument('--steps',type=int,default=1200);parser.add_argument('--eval-every',type=int,default=300);parser.add_argument('--batch-size',type=int,default=8);parser.add_argument('--lr',type=float,default=1e-4);parser.add_argument('--resume',type=Path) + opt=parser.parse_args();out=opt.output_dir;out.mkdir(exist_ok=False,parents=True) + root=Path(__file__).parent;origin=root/'runs/v101_release_20260906/parser.pt';torch.manual_seed(10106);torch.set_num_threads(4);torch.backends.cudnn.benchmark=True;torch.set_float32_matmul_precision('high') + ckpt=torch.load(origin,map_location='cpu',weights_only=False) + kwargs={k:v for k,v in ckpt['model_config'].items() if k in inspect.signature(DenseUVParserNet).parameters};kwargs['predict_head_ownership']=True + model=DenseUVParserNet(**kwargs).cuda();missing,extra=model.load_state_dict(ckpt['model'],strict=False) + if extra or any(not x.startswith('ownership_head.') for x in missing):raise ValueError((missing,extra)) + warm={k.removeprefix('accessory_head.'):v for k,v in ckpt['model'].items() if k.startswith('accessory_head.') and not k.startswith(('accessory_head.components.','accessory_head.classifier.'))} + model.ownership_head.load_state_dict(warm,strict=False) + if opt.resume:model.load_state_dict(torch.load(opt.resume,map_location='cuda',weights_only=False)['model']) + model.requires_grad_(False);model.ownership_head.requires_grad_(True) + attach_semantic_runtime(model,'siglip2','google/siglip2-base-patch16-224','cuda',local_files_only=True) + renderer=DifferentiableRenderer(ckpt['args']['mappings_dir']).cuda() + config=json.loads((root/'runs/v101_release_20260906/pipeline.json').read_text());config['head_material_refine_mode']='isolated' + split=json.loads((root/'runs/dense_uv_parser_v101_retrain_20260906/source_splits.json').read_text()) + train=OwnershipDataset(split['train'],32768,3101000);val=OwnershipDataset(split['validation'],64,29101000);test=OwnershipDataset(split['test'],128,39101000) + assert not set(split['train'])&(set(split['validation'])|set(split['test'])) + manifest={'classes':OWNERSHIP_CLASSES,'base_checkpoint':str(origin),'base_sha256':hashlib.sha256(origin.read_bytes()).hexdigest(),'options':{k:str(v) if isinstance(v,Path) else v for k,v in vars(opt).items()},'pipeline':config,'source_splits':{k:len(v) for k,v in split.items()},'real_scope':'development probes only; not training or independent test','source_sha256':{x.name:hashlib.sha256(x.read_bytes()).hexdigest() for x in root.glob('*.py')}} + write_json(out/'config.json',manifest);write_json(out/'pipeline.json',config) + snapshot=out/'source';snapshot.mkdir() + for f in root.glob('*.py'):shutil.copy2(f,snapshot/f.name) + loader=DataLoader(train,batch_size=opt.batch_size,shuffle=True,num_workers=4,pin_memory=True) + optimizer=torch.optim.AdamW(model.ownership_head.parameters(),lr=opt.lr,weight_decay=1e-4);start=time.time();step=0;epoch=0 + try: + while step=opt.steps:break + model.eval();model.ownership_head.train() + with torch.no_grad():images,target=render_ownership(batch['uv'].cuda(),batch['ownership'].cuda(),renderer,['front_left','back_left']);fg=images[:,3]>.5 + if step%2==0:images=appearance_augment(images,fg[:,None],strength=1.4) + optimizer.zero_grad(set_to_none=True) + with torch.autocast('cuda',dtype=torch.bfloat16): + logits=model.predict_ownership(images,fg);y0,y1,x0,x1=head_bounds(*target.shape[-2:]);logits=logits[:,:,y0:y1,x0:x1];truth=target[:,y0:y1,x0:x1] + ce=F.cross_entropy(logits.float(),truth,weight=images.new_tensor([1.,2.,2.,3.,4.])) + probability=logits.float().softmax(1);onehot=F.one_hot(truth,5).permute(0,3,1,2).float() + dice=1-(2*(probability*onehot).sum((2,3))+1)/(probability.sum((2,3))+onehot.sum((2,3))+1) + loss=ce+dice[:,1:].mean() + if not torch.isfinite(loss):raise RuntimeError('Non-finite loss') + loss.backward();grad=torch.nn.utils.clip_grad_norm_(model.ownership_head.parameters(),1.) + if not torch.isfinite(grad):raise RuntimeError('Non-finite gradient') + optimizer.step();step+=1 + lr=opt.lr*(.1+.9*(1+math.cos(math.pi*step/opt.steps))/2) + for g in optimizer.param_groups:g['lr']=lr + if step%25==0 or step==1: + status={'state':'training','step':step,'total_steps':opt.steps,'loss':float(loss.detach()),'elapsed_seconds':time.time()-start};write_json(out/'status.json',status);print(json.dumps(status),flush=True) + if step%opt.eval_every==0 or step==opt.steps: + metrics=evaluate(model,renderer,val);real=real_review(model,renderer,config,out,step) + record={'step':step,'validation':metrics,'real':real};write_json(out/f'evaluation_{step}.json',record) + for name,value in ckpt['model'].items(): + if not torch.equal(model.state_dict()[name].cpu(),value):raise RuntimeError('Released tensor changed: '+name) + payload={**{k:v for k,v in ckpt.items() if k not in ('model','optimizer')},'model':model.state_dict(),'model_config':{**ckpt['model_config'],'predict_head_ownership':True},'step':step,'ownership_manifest':manifest,'ownership_metrics':record} + torch.save(payload,out/f'step_{step}.pt');print('evaluation='+json.dumps(record),flush=True) + epoch+=1 + write_json(out/'heldout.json',evaluate(model,renderer,test)) + write_json(out/'status.json',{'state':'complete','step':step,'released_v101_tensors_unchanged':True,'elapsed_seconds':time.time()-start}) + except BaseException as error: + write_json(out/'status.json',{'state':'failed','step':step,'error':repr(error)});raise + + +if __name__=='__main__':main() diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0daa05f8bd4f2435fccd14287e205d57be6b958b --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/utils.py @@ -0,0 +1,4541 @@ +import math +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F + +TOOLKIT_ROOT = Path(__file__).resolve().parents[1] +WORKSPACE_ROOT = TOOLKIT_ROOT.parent +if str(WORKSPACE_ROOT) not in sys.path: + sys.path.insert(0, str(WORKSPACE_ROOT)) + +from SkingToolkit.dense_uv_parser.uv_layout import minecraft_layer_rects # noqa: E402 +from SkingToolkit.dense_uv_parser.uv_topology import ( # noqa: E402 + build_outer_uv_graph, +) + +IGNORE_INDEX = 255 +PART_CLASSES = 6 +FACE_CLASSES = 6 +LAYER_CLASSES = 2 +LAYER_FACE_CLASSES = LAYER_CLASSES * FACE_CLASSES +ROUTE_ROLE_CLASSES = 3 +ROUTE_INNER_PRIMARY = 0 +ROUTE_OUTER_PRIMARY = 1 +ROUTE_SECONDARY = 2 +UV_SIZE = 64 +SPLAT_COLOR_AGGREGATIONS = ("grid_mode", "texel_center", "exact_mode", "best") +SOLID_BACKGROUND_COLOR_TOLERANCE = 16.0 / 255.0 +SOLID_BACKGROUND_MIN_CORNER_SUPPORT = 0.5 +SECONDARY_TEXEL_MIN_PROBABILITY = 0.5 + + +def parse_views(views): + if isinstance(views, str): + return [view.strip() for view in views.split(",") if view.strip()] + return list(views) + + +def build_part_face_lookups(device=None): + inner_part = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) + inner_face = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) + outer_part = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) + outer_face = torch.full((UV_SIZE, UV_SIZE), IGNORE_INDEX, dtype=torch.long, device=device) + + rects = minecraft_layer_rects(is_slim=False) + for rect_index, (inner_x, inner_y, width, height, decor_dx, decor_dy) in enumerate(rects): + part = rect_index // FACE_CLASSES + face = rect_index % FACE_CLASSES + inner_part[inner_y : inner_y + height, inner_x : inner_x + width] = part + inner_face[inner_y : inner_y + height, inner_x : inner_x + width] = face + + outer_x = inner_x + decor_dx + outer_y = inner_y + decor_dy + outer_part[outer_y : outer_y + height, outer_x : outer_x + width] = part + outer_face[outer_y : outer_y + height, outer_x : outer_x + width] = face + + return { + "inner_part": inner_part, + "inner_face": inner_face, + "outer_part": outer_part, + "outer_face": outer_face, + } + + +def _grid_to_uv01(grid): + return ((grid + 1.0) * 0.5).clamp(0.0, 1.0) + + +def _grid_to_xy(grid): + coords = (_grid_to_uv01(grid) * (UV_SIZE - 1)).round().long() + x = coords[..., 0].clamp(0, UV_SIZE - 1) + y = coords[..., 1].clamp(0, UV_SIZE - 1) + return x, y + + +def _renderer_surface_cache(renderer): + cache = getattr(renderer, "_dense_uv_surface_cache", None) + if cache is None: + cache = {} + setattr(renderer, "_dense_uv_surface_cache", cache) + return cache + + +def _surface_metadata(grids, is_outer, lookups): + uv_px = _grid_to_uv01(grids) * (UV_SIZE - 1) + rounded_uv_px = uv_px.round() + x = rounded_uv_px[..., 0].long().clamp(0, UV_SIZE - 1) + y = rounded_uv_px[..., 1].long().clamp(0, UV_SIZE - 1) + center_distance_sq = (uv_px - rounded_uv_px).square().sum(dim=-1) + # A bilinear renderer is exact at integer UV centers and increasingly + # mixes neighboring texels toward a cell boundary. Keep a normalized + # quality score for inverse color sampling; geometry/UV ownership still + # uses the same rounded integer coordinates as before. + texel_center_score = (1.0 - 2.0 * center_distance_sq).clamp(0.0, 1.0) + inner_part = lookups["inner_part"][y, x] + inner_face = lookups["inner_face"][y, x] + outer_part = lookups["outer_part"][y, x] + outer_face = lookups["outer_face"][y, x] + return { + "flat_uv": y * UV_SIZE + x, + "layer": is_outer.long(), + "part": torch.where(is_outer, outer_part, inner_part), + "face": torch.where(is_outer, outer_face, inner_face), + "outer_part": outer_part, + "texel_center_score": texel_center_score, + } + + +def build_static_surface_routing(renderer, view, device): + """Build fixed UV candidates matching the renderer's direct/composite geometry paths.""" + device = torch.device(device) + cache = _renderer_surface_cache(renderer) + cache_key = (view, device.type, device.index) + if cache_key in cache: + return cache[cache_key] + + inner_grid = getattr(renderer, f"{view}_inner_grid").to(device=device) + outer_grid = getattr(renderer, f"{view}_outer_grid").to(device=device) + inner_mask = getattr(renderer, f"{view}_inner_mask").to(device=device).bool() + outer_mask = getattr(renderer, f"{view}_outer_mask").to(device=device).bool() + H, W = inner_mask.shape + lookups = build_part_face_lookups(device=device) + + grids = [inner_grid, outer_grid] + masks = [inner_mask, outer_mask] + is_outer = [torch.zeros((H, W), dtype=torch.bool, device=device), torch.ones((H, W), dtype=torch.bool, device=device)] + composite_count = 0 + geometry_count = 0 + + composite_grid_name = f"{view}_composite_grid_layers" + composite_decor_name = f"{view}_composite_is_decor_layers" + if hasattr(renderer, composite_grid_name) and hasattr(renderer, composite_decor_name): + composite_grids = getattr(renderer, composite_grid_name).to(device=device) + composite_masks = getattr(renderer, f"{view}_composite_mask_layers").to(device=device).bool() + composite_decor = getattr(renderer, composite_decor_name).to(device=device).bool() + composite_count = composite_grids.shape[0] + grids.extend(composite_grids.unbind(0)) + masks.extend(composite_masks.unbind(0)) + is_outer.extend(composite_decor.unbind(0)) + + geometry_grid_name = f"{view}_geometry_grid_layers" + if hasattr(renderer, geometry_grid_name): + geometry_grids = getattr(renderer, geometry_grid_name).to(device=device) + geometry_masks = getattr(renderer, f"{view}_geometry_mask_layers").to(device=device).bool() + geometry_order = getattr(renderer, f"{view}_geometry_sort_indices").to(device=device).long() + geometry_count = geometry_grids.shape[0] + if geometry_count > 0: + geometry_order = geometry_order.clamp(0, geometry_count - 1) + sorted_grids = torch.gather( + geometry_grids, + 0, + geometry_order.unsqueeze(-1).expand(-1, -1, -1, 2), + ) + sorted_masks = torch.gather(geometry_masks, 0, geometry_order) + geometry_metadata = _surface_metadata( + sorted_grids, + torch.zeros_like(sorted_masks), + lookups, + ) + geometry_is_outer = geometry_metadata["outer_part"] != IGNORE_INDEX + grids.extend(sorted_grids.unbind(0)) + masks.extend(sorted_masks.unbind(0)) + is_outer.extend(geometry_is_outer.unbind(0)) + + surface_grids = torch.stack(grids, dim=0) + surface_masks = torch.stack(masks, dim=0) + surface_is_outer = torch.stack(is_outer, dim=0) + metadata = _surface_metadata(surface_grids, surface_is_outer, lookups) + metadata["layer_face"] = combine_layer_face(metadata["layer"], metadata["face"]) + static = { + "grids": surface_grids, + "masks": surface_masks, + "flat_uv": metadata["flat_uv"], + "layer": metadata["layer"], + "part": metadata["part"], + "face": metadata["face"], + "layer_face": metadata["layer_face"], + "texel_center_score": metadata["texel_center_score"], + "direct_count": 2, + "composite_count": composite_count, + "geometry_count": geometry_count, + } + cache[cache_key] = static + return static + + +def surface_class_count(renderer, views): + device = renderer.bg_color.device + return max(build_static_surface_routing(renderer, view, device)["grids"].shape[0] for view in parse_views(views)) + + +def _select_static_surface(static, surface): + B, H, W = surface.shape + surface_count = static["grids"].shape[0] + in_range = (surface >= 0) & (surface < surface_count) + safe_surface = surface.clamp(0, max(surface_count - 1, 0)) + + selected = {} + for name in ( + "masks", + "flat_uv", + "layer", + "part", + "face", + "texel_center_score", + ): + values = static[name].unsqueeze(0).expand(B, -1, -1, -1) + selected[name] = values.gather(1, safe_surface.unsqueeze(1)).squeeze(1) + selected["valid"] = in_range & selected["masks"] + return selected + + +def _sample_surface_alpha(skins, static): + B = skins.shape[0] + surface_count, H, W, _ = static["grids"].shape + grids = static["grids"].unsqueeze(0).expand(B, -1, -1, -1, -1).reshape(B * surface_count, H, W, 2) + alpha = skins[:, 3:4].unsqueeze(1).expand(-1, surface_count, -1, -1, -1) + alpha = alpha.reshape(B * surface_count, 1, skins.shape[-2], skins.shape[-1]) + sampled = F.grid_sample(alpha, grids, mode="bilinear", padding_mode="zeros", align_corners=True) + return sampled.reshape(B, surface_count, H, W) * static["masks"].unsqueeze(0) + + +def classify_route_role(static, selected_layer, selected_flat_uv, valid): + """Classify visible pixels as direct inner/outer surfaces or deeper backfaces.""" + B, H, W = selected_layer.shape + safe_layer = selected_layer.clamp(0, LAYER_CLASSES - 1) + direct_masks = static["masks"][:LAYER_CLASSES].unsqueeze(0).expand(B, -1, H, W) + direct_uv = static["flat_uv"][:LAYER_CLASSES].unsqueeze(0).expand(B, -1, H, W) + direct_valid = direct_masks.gather(1, safe_layer.unsqueeze(1)).squeeze(1) + expected_uv = direct_uv.gather(1, safe_layer.unsqueeze(1)).squeeze(1) + primary = valid & direct_valid & (selected_flat_uv == expected_uv) + route_role = torch.full_like(selected_layer, IGNORE_INDEX) + route_role[primary] = selected_layer[primary] + route_role[valid & ~primary] = ROUTE_SECONDARY + return route_role + + +def build_dense_parser_batch(skins, renderer, view, alpha_threshold=0.5): + """Render skins and label the exact fixed surface that produced every pixel.""" + device = skins.device + dtype = skins.dtype + B = skins.shape[0] + + rendered = renderer.forward_view(skins, view) + static = build_static_surface_routing(renderer, view, device) + _, _, H, W = rendered.shape + sampled_alpha = _sample_surface_alpha(skins, static) + + surface = torch.full((B, H, W), IGNORE_INDEX, dtype=torch.long, device=device) + direct_visible = sampled_alpha[:, :2] > alpha_threshold + inner_visible = direct_visible[:, 0] & ~direct_visible[:, 1] + outer_visible = direct_visible[:, 1] + surface[inner_visible] = 0 + surface[outer_visible] = 1 + + composite_start = static["direct_count"] + composite_count = static["composite_count"] + composite_visible = None + if composite_count > 0: + composite_visible = sampled_alpha[:, composite_start : composite_start + composite_count] > alpha_threshold + composite_indices = torch.arange(composite_count, device=device).view(1, composite_count, 1, 1) + first_composite = torch.where( + composite_visible, + composite_indices, + torch.full_like(composite_indices, composite_count), + ).amin(dim=1) + has_composite = first_composite < composite_count + composite_surface = composite_start + first_composite.clamp(max=composite_count - 1) + composite_route = _select_static_surface(static, composite_surface) + trust_composite = has_composite & ( + (composite_route["layer"] == 1) | (first_composite <= 1) + ) + surface[trust_composite] = composite_surface[trust_composite] + + geometry_start = composite_start + composite_count + geometry_count = static["geometry_count"] + if geometry_count > 0: + geometry_visible = sampled_alpha[:, geometry_start : geometry_start + geometry_count] > alpha_threshold + geometry_indices = torch.arange(geometry_count, device=device).view(1, geometry_count, 1, 1) + last_geometry = torch.where( + geometry_visible, + geometry_indices, + torch.full_like(geometry_indices, -1), + ).amax(dim=1) + has_geometry = last_geometry >= 0 + geometry_surface = geometry_start + last_geometry.clamp_min(0) + front_decor = static["masks"][composite_start] & (static["layer"][composite_start] == 1) + geometry_fallback = ( + front_decor.unsqueeze(0) + & ~composite_visible[:, 0] + & has_geometry + ) + surface[geometry_fallback] = geometry_surface[geometry_fallback] + + selected = _select_static_surface(static, surface) + valid = selected["valid"] + layer = torch.where(valid, selected["layer"], torch.full_like(surface, IGNORE_INDEX)) + route_role = classify_route_role(static, selected["layer"], selected["flat_uv"], valid) + part = torch.where(valid, selected["part"], torch.full_like(surface, IGNORE_INDEX)) + face = torch.where(valid, selected["face"], torch.full_like(surface, IGNORE_INDEX)) + uv = flat_uv_to_uv01(selected["flat_uv"], dtype).masked_fill(~valid.unsqueeze(1), 0.0) + + targets = { + "foreground": valid.unsqueeze(1).to(dtype=dtype), + "layer": layer, + "route_role": route_role, + "part": part, + "face": face, + "surface": torch.where(valid, surface, torch.full_like(surface, IGNORE_INDEX)), + "uv": uv, + } + return rendered, targets + + +def randomize_render_background(rendered, probability=0.9, bg_color=(128, 128, 128)): + """Replace rendered backgrounds while preserving the skin pixels and parser targets. + + The renderer produces RGBA images composited over a fixed gray background. The + parser always receives RGB-style inputs (alpha fixed to one), with a random + solid-color background for the requested fraction of samples. + """ + if rendered.dim() != 4 or rendered.shape[1] != 4: + raise ValueError(f"Expected RGBA render tensor as NCHW, got {tuple(rendered.shape)}.") + + B, _, H, W = rendered.shape + device = rendered.device + dtype = rendered.dtype + alpha = rendered[:, 3:4].clamp(0.0, 1.0) + source_bg = rendered.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 + probability = max(0.0, min(float(probability), 1.0)) + if probability > 0: + active = torch.rand(B, device=device) < probability + random_background = torch.rand(B, 3, 1, 1, device=device, dtype=dtype).expand(-1, -1, H, W) + background = torch.where(active.view(B, 1, 1, 1), random_background, source_bg.expand(B, -1, H, W)) + else: + background = source_bg.expand(B, -1, H, W) + foreground_rgb = (rendered[:, :3] - (1.0 - alpha) * source_bg) / alpha.clamp_min(1e-4) + composited_rgb = (alpha * foreground_rgb + (1.0 - alpha) * background).clamp(0.0, 1.0) + + return torch.cat([composited_rgb, torch.ones_like(alpha)], dim=1) + + +def estimate_solid_background_color( + rendered, + color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, + min_corner_support=SOLID_BACKGROUND_MIN_CORNER_SUPPORT, +): + if rendered.dim() != 4 or rendered.shape[1] < 3: + raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") + if not 0.0 <= color_tolerance <= 1.0: + raise ValueError("color_tolerance must be in [0, 1].") + if not 0.0 <= min_corner_support <= 1.0: + raise ValueError("min_corner_support must be in [0, 1].") + + rgb = rendered[:, :3].float() + batch, _, height, width = rgb.shape + patch = max(2, min(32, min(height, width) // 16)) + corners = torch.cat( + [ + rgb[:, :, :patch, :patch].flatten(2), + rgb[:, :, :patch, -patch:].flatten(2), + rgb[:, :, -patch:, :patch].flatten(2), + rgb[:, :, -patch:, -patch:].flatten(2), + ], + dim=2, + ) + background_rgb = corners.median(dim=2).values.view(batch, 3, 1, 1) + corner_distance = (corners - background_rgb.flatten(2)).abs().amax(dim=1) + solid_background = (corner_distance <= color_tolerance).float().mean(dim=1) >= min_corner_support + return background_rgb, solid_background + + +def estimate_solid_background_foreground( + rendered, + color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, + min_corner_support=SOLID_BACKGROUND_MIN_CORNER_SUPPORT, +): + """Find non-background pixels without deleting matching colors enclosed by the subject.""" + if min(rendered.shape[-2:]) < 8: + return torch.ones( + rendered.shape[0], + *rendered.shape[-2:], + dtype=torch.bool, + device=rendered.device, + ) + + rgb = rendered[:, :3].float() + batch, _, height, width = rgb.shape + background_rgb, solid_background = estimate_solid_background_color( + rendered, + color_tolerance=color_tolerance, + min_corner_support=min_corner_support, + ) + + background_candidate = (rgb - background_rgb).abs().amax(dim=1) <= color_tolerance + candidate_u8 = background_candidate.to(torch.uint8) + connected_to_border = ( + candidate_u8.cumprod(dim=2).bool() + | candidate_u8.flip(2).cumprod(dim=2).flip(2).bool() + | candidate_u8.cumprod(dim=1).bool() + | candidate_u8.flip(1).cumprod(dim=1).flip(1).bool() + ) + connected_to_border &= solid_background.view(batch, 1, 1) + max_steps = max(height, width) + for _ in range(max_steps): + expanded = F.max_pool2d( + connected_to_border.unsqueeze(1).float(), 3, stride=1, padding=1 + )[:, 0] > 0.0 + updated = connected_to_border | (background_candidate & expanded) + if torch.equal(updated, connected_to_border): + break + connected_to_border = updated + return ~connected_to_border + + +def estimate_top_left_flood_foreground(rendered, color_tolerance=8.0 / 255.0): + """Remove only seed-colored pixels connected to the top-left image pixel. + + Each image in the batch owns its own RGB seed. Pixels with the same color + elsewhere inside the character remain foreground unless a path of + seed-colored pixels connects them to the top-left corner. + """ + if rendered.dim() != 4 or rendered.shape[1] < 3: + raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") + if not 0.0 <= color_tolerance <= 1.0: + raise ValueError("color_tolerance must be in [0, 1].") + + rgb = rendered[:, :3].float() + seed_rgb = rgb[:, :, :1, :1] + background_candidate = ( + (rgb - seed_rgb).abs().amax(dim=1) <= color_tolerance + ) + connected = torch.zeros_like(background_candidate) + connected[:, 0, 0] = background_candidate[:, 0, 0] + height, width = rendered.shape[-2:] + for _ in range(height + width): + # Four-connected expansion cannot leak diagonally through a one-pixel + # corner where two foreground blocks merely touch. + expanded = connected.clone() + expanded[:, 1:, :] |= connected[:, :-1, :] + expanded[:, :-1, :] |= connected[:, 1:, :] + expanded[:, :, 1:] |= connected[:, :, :-1] + expanded[:, :, :-1] |= connected[:, :, 1:] + updated = connected | (background_candidate & expanded) + if torch.equal(updated, connected): + break + connected = updated + return ~connected + + +def build_color_sampling_support( + rendered, + observed_foreground, + background_rgb, + background_tolerance=None, + foreground_inset=1, +): + """Separate route occupancy from source pixels that are safe for RGB pickup. + + Flood fill intentionally keeps colors that merely resemble the background + when they are not connected to the image border. That is correct for + occupancy, but an isolated background pocket or antialiased silhouette + fringe must not beat a real character sample during inverse UV splatting. + Only background-like *boundary* candidates are rejected. Interior pixels + remain valid even when a skin intentionally uses the background color. + """ + if rendered.dim() != 4 or rendered.shape[1] < 3: + raise ValueError(f"Expected NCHW RGB(A), got {tuple(rendered.shape)}.") + expected = rendered.shape[:1] + rendered.shape[-2:] + if observed_foreground.shape != expected: + raise ValueError( + f"Expected observed_foreground shape {expected}, got " + f"{tuple(observed_foreground.shape)}." + ) + if background_rgb.shape not in ((rendered.shape[0], 3), (rendered.shape[0], 3, 1, 1)): + raise ValueError( + "background_rgb must have shape (N, 3) or (N, 3, 1, 1)." + ) + if background_tolerance is not None and not 0.0 <= background_tolerance <= 1.0: + raise ValueError("background_tolerance must be in [0, 1].") + if foreground_inset < 0: + raise ValueError("foreground_inset must be non-negative.") + + observed = observed_foreground.to(device=rendered.device, dtype=torch.bool) + interior = observed + for _ in range(int(foreground_inset)): + adjacent_background = F.max_pool2d( + (~interior).unsqueeze(1).float(), + kernel_size=3, + stride=1, + padding=1, + )[:, 0] > 0.0 + interior = interior & ~adjacent_background + + background = background_rgb.to( + device=rendered.device, dtype=rendered.dtype + ).reshape(rendered.shape[0], 3, 1, 1) + background_like = ( + (rendered[:, :3] - background).abs().amax(dim=1) + <= float(background_tolerance) + if background_tolerance is not None + else torch.zeros_like(observed) + ) + boundary = observed & ~interior + rejected = boundary & background_like + return { + "valid": observed & ~rejected, + "interior": interior, + "background_like": background_like, + "rejected": rejected, + } + + +def affine_to_canonical_grid(affine, output_shape): + """Build a grid that samples an augmented render at canonical coordinates.""" + if affine.dim() != 2 or affine.shape[1] != 3: + raise ValueError(f"Expected affine shape (N, 3), got {tuple(affine.shape)}.") + if affine.shape[0] != output_shape[0]: + raise ValueError( + f"Affine batch size {affine.shape[0]} does not match output batch size {output_shape[0]}." + ) + + scale = affine[:, 2].exp() + theta = affine.new_zeros(affine.shape[0], 2, 3) + theta[:, 0, 0] = scale + theta[:, 1, 1] = scale + theta[:, 0, 2] = affine[:, 0] + theta[:, 1, 2] = affine[:, 1] + return F.affine_grid(theta, output_shape, align_corners=False) + + +def canonicalize_tensor(tensor, affine, mode="bilinear"): + """Undo the parser's predicted global translation/scale for an NCHW tensor.""" + affine = affine.to(device=tensor.device, dtype=tensor.dtype) + grid = affine_to_canonical_grid(affine, tensor.shape) + return F.grid_sample( + tensor, + grid, + mode=mode, + padding_mode="zeros", + align_corners=False, + ) + + +def canonicalize_foreground_mask(mask, affine, min_coverage=1e-3): + """Warp a foreground mask without dropping thin silhouette coverage. + + Nearest-neighbour sampling can erase a one-pixel hat brim or outer-layer + protrusion after even a small predicted affine scale. Bilinear coverage + with a low positive threshold keeps every output pixel touched by the + transformed source mask. Later color/background and minimum-source filters + remain responsible for rejecting unsupported edge colors. + """ + if mask.dim() == 3: + mask = mask.unsqueeze(1) + if mask.dim() != 4 or mask.shape[1] != 1: + raise ValueError( + "Foreground mask must have shape NHW or N1HW, got " + f"{tuple(mask.shape)}." + ) + if not 0.0 < float(min_coverage) <= 1.0: + raise ValueError("min_coverage must be in (0, 1].") + coverage = canonicalize_tensor( + mask.to(device=affine.device, dtype=torch.float32), + affine, + mode="bilinear", + ) + return coverage[:, 0] >= float(min_coverage) + + +def canonicalize_index_tensor(index, affine): + """Undo a global transform while preserving IGNORE_INDEX outside valid labels.""" + valid = (index != IGNORE_INDEX).unsqueeze(1).float() + values = index.masked_fill(index == IGNORE_INDEX, 0).float().unsqueeze(1) + affine = affine.to(device=values.device, dtype=values.dtype) + grid = affine_to_canonical_grid(affine, values.shape) + sampled_valid = F.grid_sample( + valid, + grid, + mode="nearest", + padding_mode="zeros", + align_corners=False, + ).squeeze(1) > 0.5 + sampled_values = F.grid_sample( + values, + grid, + mode="nearest", + padding_mode="zeros", + align_corners=False, + ).squeeze(1).round().long() + sampled_values[~sampled_valid] = IGNORE_INDEX + return sampled_values + + +def canonicalize_dense_targets(targets): + """Return parser targets in the fixed renderer coordinate system.""" + if "affine" not in targets: + raise ValueError("Global-affine targets require an 'affine' tensor.") + affine = targets["affine"] + canonical = { + "foreground": canonicalize_tensor(targets["foreground"], affine, mode="nearest"), + "layer": canonicalize_index_tensor(targets["layer"], affine), + "route_role": canonicalize_index_tensor(targets["route_role"], affine), + "part": canonicalize_index_tensor(targets["part"], affine), + "face": canonicalize_index_tensor(targets["face"], affine), + "surface": canonicalize_index_tensor(targets["surface"], affine), + "uv": canonicalize_tensor(targets["uv"], affine, mode="nearest"), + "affine": affine, + } + valid = (canonical["layer"] != IGNORE_INDEX).unsqueeze(1) + canonical["foreground"] = canonical["foreground"] * valid.to(canonical["foreground"].dtype) + canonical["uv"] = torch.where(valid, canonical["uv"], torch.zeros_like(canonical["uv"])) + return canonical + + +def canonicalize_parser_outputs(outputs): + """Warp dense parser logits into the fixed renderer coordinate system.""" + if "affine" not in outputs: + return outputs + affine = outputs["affine"] + canonical = {} + for name, value in outputs.items(): + if ( + name in ( + "affine", + "route_role_prior_raw", + "outer_uv_occupancy_logits", + "outer_uv_occupancy_support", + ) + or name == "outer_uv_global_context" + or not torch.is_tensor(value) + or value.dim() != 4 + ): + canonical[name] = value + else: + canonical[name] = canonicalize_tensor(value, affine, mode="bilinear") + return canonical + + +def canonicalize_parser_render(rendered, outputs, mode="nearest", fill_color=None): + """Undo the global transform for colors without blending Minecraft texels.""" + if "affine" not in outputs: + return rendered + if fill_color is None: + return canonicalize_tensor(rendered, outputs["affine"], mode=mode) + if fill_color.dim() == 2: + fill_color = fill_color.unsqueeze(-1).unsqueeze(-1) + if fill_color.shape != (rendered.shape[0], rendered.shape[1], 1, 1): + raise ValueError( + "fill_color must have shape " + f"{(rendered.shape[0], rendered.shape[1], 1, 1)}, got {tuple(fill_color.shape)}." + ) + fill_color = fill_color.to(device=rendered.device, dtype=rendered.dtype) + return canonicalize_tensor( + rendered - fill_color, outputs["affine"], mode=mode + ) + fill_color + + +def _mask_moments(weights): + weights = weights.float() + N, H, W = weights.shape + y = torch.arange(H, device=weights.device, dtype=weights.dtype).view(1, H, 1) + x = torch.arange(W, device=weights.device, dtype=weights.dtype).view(1, 1, W) + total = weights.sum(dim=(1, 2)).clamp_min(1e-6) + mean_x = (weights * x).sum(dim=(1, 2)) / total + mean_y = (weights * y).sum(dim=(1, 2)) / total + var_x = (weights * (x - mean_x.view(N, 1, 1)).square()).sum(dim=(1, 2)) / total + var_y = (weights * (y - mean_y.view(N, 1, 1)).square()).sum(dim=(1, 2)) / total + return total, mean_x, mean_y, var_x.clamp_min(1e-6).sqrt(), var_y.clamp_min(1e-6).sqrt() + + +def _soft_mask_dice(probability, target, support): + probability = probability.float() * support.float() + target = target.float() + intersection = (probability * target).sum(dim=(1, 2, 3)) + denominator = probability.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) + return (2.0 * intersection + 1e-6) / (denominator + 1e-6) + + +def refine_parser_affine( + outputs, + renderer, + views, + translation_radius_px=8.0, + scale_radius=0.0, + translation_step_px=0.5, + observed_foreground=None, +): + """Refine parser affine predictions against the fixed inner-layer silhouette.""" + if "affine" not in outputs or "foreground" not in outputs: + return outputs.get("affine"), None + if translation_radius_px < 0 or scale_radius < 0: + raise ValueError("Affine refinement radii must be non-negative.") + if translation_step_px <= 0: + raise ValueError("translation_step_px must be positive.") + + views = parse_views(views) + affine = outputs["affine"] + if observed_foreground is None: + foreground_prob = torch.sigmoid(outputs["foreground"]) + else: + if observed_foreground.shape != outputs["foreground"].shape[:1] + outputs["foreground"].shape[-2:]: + raise ValueError( + "observed_foreground must match parser spatial output shape, got " + f"{tuple(observed_foreground.shape)}." + ) + foreground_prob = observed_foreground.to( + device=outputs["foreground"].device, + dtype=outputs["foreground"].dtype, + ).unsqueeze(1) + N, _, H, W = foreground_prob.shape + if not views or N % len(views) != 0: + raise ValueError(f"N={N} must be divisible by the number of views ({len(views)}).") + + target = foreground_prob.new_zeros(N, 1, H, W) + for view_index, view in enumerate(views): + static = build_static_surface_routing(renderer, view, foreground_prob.device) + inner_mask = static["masks"][0] + if inner_mask.shape != (H, W): + raise ValueError( + f"View {view!r} mapping shape {tuple(inner_mask.shape)} does not match parser shape {(H, W)}." + ) + target[view_index::len(views), 0] = inner_mask.to(dtype=target.dtype) + + support_radius = max(math.ceil(translation_radius_px) + 2, 2) + support = F.max_pool2d( + target, + kernel_size=support_radius * 2 + 1, + stride=1, + padding=support_radius, + ).clamp(0.0, 1.0) + + base_canonical = canonicalize_tensor(foreground_prob, affine, mode="bilinear") + base_score = _soft_mask_dice(base_canonical, target, support) + target_stats = _mask_moments(target[:, 0]) + base_stats = _mask_moments(base_canonical[:, 0] * support[:, 0]) + target_total, _, _, target_std_x, target_std_y = target_stats + base_total, _, _, base_std_x, base_std_y = base_stats + enough_foreground = base_total > target_total * 0.25 + + if scale_radius > 0: + scale_delta = 0.5 * ( + (base_std_x / target_std_x).clamp_min(1e-6).log() + + (base_std_y / target_std_y).clamp_min(1e-6).log() + ) + scale_delta = scale_delta.clamp(-scale_radius, scale_radius) + scale_affine = affine.clone() + scale_affine[:, 2] = scale_affine[:, 2] + scale_delta.to(dtype=scale_affine.dtype) + scale_canonical = canonicalize_tensor(foreground_prob, scale_affine, mode="bilinear") + else: + scale_delta = torch.zeros_like(base_score) + scale_affine = affine + scale_canonical = base_canonical + _, pred_mean_x, pred_mean_y, _, _ = _mask_moments(scale_canonical[:, 0] * support[:, 0]) + _, target_mean_x, target_mean_y, _, _ = target_stats + delta_x_px = (pred_mean_x - target_mean_x).clamp(-translation_radius_px, translation_radius_px) + delta_y_px = (pred_mean_y - target_mean_y).clamp(-translation_radius_px, translation_radius_px) + delta_x_px = (delta_x_px / translation_step_px).round() * translation_step_px + delta_y_px = (delta_y_px / translation_step_px).round() * translation_step_px + + candidate_affine = scale_affine.clone() + candidate_affine[:, 0] = candidate_affine[:, 0] + (2.0 * delta_x_px / W).to(candidate_affine.dtype) + candidate_affine[:, 1] = candidate_affine[:, 1] + (2.0 * delta_y_px / H).to(candidate_affine.dtype) + candidate_canonical = canonicalize_tensor(foreground_prob, candidate_affine, mode="bilinear") + candidate_score = _soft_mask_dice(candidate_canonical, target, support) + accepted = enough_foreground & (candidate_score > base_score + 1e-5) + refined_affine = torch.where(accepted.unsqueeze(1), candidate_affine, affine) + + accepted_float = accepted.to(dtype=delta_x_px.dtype) + details = { + "accepted": accepted, + "translation_px": torch.stack( + [delta_x_px * accepted_float, delta_y_px * accepted_float], + dim=1, + ), + "scale_percent": (scale_delta.exp() - 1.0) * 100.0 * accepted_float, + "score_before": base_score, + "score_after": torch.where(accepted, candidate_score, base_score), + } + return refined_affine, details + + +def _routing_from_affine_outputs(renderer, views, outputs, fg_threshold=0.5, semantic_gate=True): + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required for deterministic UV routing.") + if "surface" not in outputs: + raise ValueError("Affine parser outputs must include the static surface classifier.") + foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] + surface_prob = torch.softmax(outputs["surface"], dim=1) + layer_prob = torch.softmax(outputs["layer"], dim=1) + N, _, H, W = surface_prob.shape + if N % len(views) != 0: + raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") + + fg = torch.zeros_like(foreground_prob, dtype=torch.bool) + layer = torch.zeros((N, H, W), dtype=torch.long, device=surface_prob.device) + flat_uv = torch.zeros_like(layer) + surface = torch.full_like(layer, IGNORE_INDEX) + confidence = torch.zeros_like(foreground_prob) + confidence_margin = torch.zeros_like(foreground_prob) + confidence_margin_ratio = torch.zeros_like(foreground_prob) + texel_center_score = torch.zeros_like(foreground_prob) + semantic_fallback = torch.zeros_like(foreground_prob, dtype=torch.bool) + expected_part = torch.full_like(layer, IGNORE_INDEX) + expected_face = torch.full_like(layer, IGNORE_INDEX) + + for view_index, view in enumerate(views): + selection = slice(view_index, N, len(views)) + static = build_static_surface_routing(renderer, view, surface_prob.device) + if static["masks"].shape[-2:] != (H, W): + raise ValueError( + f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} does not match parser shape {(H, W)}." + ) + view_surface_prob = surface_prob[selection] + view_batch = view_surface_prob.shape[0] + surface_count = static["masks"].shape[0] + if view_surface_prob.shape[1] < surface_count: + raise ValueError( + f"Parser has {view_surface_prob.shape[1]} surface classes, but view {view!r} requires {surface_count}." + ) + + candidate_score = view_surface_prob[:, :surface_count] + physical_valid = static["masks"].unsqueeze(0).expand(view_batch, -1, -1, -1) + candidate_valid = physical_valid.clone() + semantic_score = torch.ones_like(candidate_score) + + # Re-rank only surfaces that physically exist at this screen pixel. This + # converts the semantic heads from a rejection gate into useful routing + # evidence and avoids holes caused by a globally invalid surface argmax. + semantic_names = ( + ("part", "layer_face") + if "layer_face" in outputs + else ("layer", "part", "face") + ) + for name in semantic_names: + if name not in outputs: + continue + probabilities = torch.softmax(outputs[name][selection], dim=1) + expected = static[name] + known_semantic = expected != IGNORE_INDEX + expected_safe = expected.clamp(0, probabilities.shape[1] - 1) + expected_probability = probabilities.gather( + 1, + expected_safe.unsqueeze(0).expand(view_batch, -1, -1, -1), + ) + semantic_score = semantic_score * torch.where( + known_semantic.unsqueeze(0), + expected_probability, + torch.ones_like(expected_probability), + ) + if semantic_gate: + candidate_valid = candidate_valid & ( + ~known_semantic.unsqueeze(0) + | (probabilities.argmax(dim=1).unsqueeze(1) == expected_safe.unsqueeze(0)) + ) + + candidate_score = candidate_score * semantic_score.clamp_min(1e-12).sqrt() + if "uv_x" in outputs and "uv_y" in outputs: + uv_x_prob = torch.softmax(outputs["uv_x"][selection], dim=1) + uv_y_prob = torch.softmax(outputs["uv_y"][selection], dim=1) + candidate_x = static["flat_uv"].remainder(UV_SIZE) + candidate_y = static["flat_uv"].div(UV_SIZE, rounding_mode="floor") + candidate_x_prob = uv_x_prob.gather( + 1, + candidate_x.unsqueeze(0).expand(view_batch, -1, -1, -1), + ) + candidate_y_prob = uv_y_prob.gather( + 1, + candidate_y.unsqueeze(0).expand(view_batch, -1, -1, -1), + ) + candidate_score = candidate_score * ( + candidate_x_prob * candidate_y_prob + ).clamp_min(1e-12).sqrt() + gated_score = candidate_score.masked_fill(~candidate_valid, -1.0) + has_gated_candidate = candidate_valid.any(dim=1, keepdim=True) + if semantic_gate: + fallback_score = candidate_score.masked_fill(~physical_valid, -1.0) + candidate_score = torch.where(has_gated_candidate, gated_score, fallback_score) + semantic_fallback[selection] = ( + ~has_gated_candidate.squeeze(1) & physical_valid.any(dim=1) + ) + else: + candidate_score = gated_score + best_score, selected_surface = candidate_score.max(dim=1) + if surface_count > 1: + top_scores = candidate_score.topk(k=2, dim=1).values + route_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) + route_margin_ratio = ( + route_margin / top_scores[:, 0].abs().clamp_min(1e-8) + ).clamp(0.0, 1.0) + else: + route_margin = best_score.clamp_min(0.0) + route_margin_ratio = (best_score > 0.0).to(best_score.dtype) + + routed = _select_static_surface(static, selected_surface) + selected_layer = routed["layer"] + selected_part = routed["part"] + selected_face = routed["face"] + has_candidate = best_score >= 0.0 + base_silhouette = static["masks"][0].unsqueeze(0).expand(view_batch, -1, -1) + routed_fg = ( + ((foreground_prob[selection] > fg_threshold) | base_silhouette) + & has_candidate + & routed["valid"] + ) + fg[selection] = routed_fg + layer[selection] = selected_layer + flat_uv[selection] = routed["flat_uv"] + surface[selection] = torch.where( + has_candidate, + selected_surface, + torch.full_like(selected_surface, IGNORE_INDEX), + ) + confidence[selection] = foreground_prob[selection] * best_score.clamp_min(0.0) + confidence_margin[selection] = route_margin + confidence_margin_ratio[selection] = route_margin_ratio + expected_part[selection] = selected_part + expected_face[selection] = selected_face + texel_center_score[selection] = routed["texel_center_score"] + + return { + "foreground": fg, + "layer": layer, + "flat_uv": flat_uv, + "surface": surface, + "confidence": confidence, + "confidence_margin": confidence_margin, + "confidence_margin_ratio": confidence_margin_ratio, + "semantic_fallback": semantic_fallback, + "part": expected_part, + "face": expected_face, + "texel_center_score": texel_center_score, + } + + +def _aggregate_role_probabilities_by_direct_texel( + role_prob, + static, + foreground_prob=None, + center_power=2.0, +): + """Center-weight route probabilities within each projected Minecraft texel.""" + role_prob = role_prob.float() + batch, role_count, height, width = role_prob.shape + if foreground_prob is not None and foreground_prob.shape != ( + batch, + height, + width, + ): + raise ValueError( + "foreground_prob must match the route map, got " + f"{tuple(foreground_prob.shape)} and {(batch, height, width)}." + ) + if center_power < 0: + raise ValueError("center_power must be non-negative.") + aggregated = [] + for layer_index in range(LAYER_CLASSES): + valid = static["masks"][layer_index].reshape(1, 1, -1) + flat_uv = static["flat_uv"][layer_index].reshape(1, 1, -1) + center_weight = static["texel_center_score"][layer_index].float() + center_weight = center_weight.clamp(0.0, 1.0).pow(center_power) + pixel_weight = center_weight.reshape(1, 1, -1) * valid + if foreground_prob is not None: + pixel_weight = ( + pixel_weight + * foreground_prob.float().clamp(0.0, 1.0).reshape(batch, 1, -1) + ) + pixel_weight = pixel_weight.expand(batch, 1, -1) + sums = role_prob.new_zeros(batch, role_count, UV_SIZE * UV_SIZE) + sums.scatter_add_( + 2, + flat_uv.expand(batch, role_count, -1), + role_prob.flatten(2) * pixel_weight, + ) + counts = role_prob.new_zeros(batch, 1, UV_SIZE * UV_SIZE) + counts.scatter_add_( + 2, + flat_uv.expand(batch, 1, -1), + pixel_weight, + ) + means = sums / counts.clamp_min(1e-6) + aggregated.append( + means.gather(2, flat_uv.expand(batch, role_count, -1)).reshape( + batch, role_count, height, width + ) + ) + return torch.stack(aggregated, dim=1) + + +def _project_outer_uv_occupancy_to_view( + occupancy_probability, + static, + view_batch, +): + """Gather a grouped 64x64 outer-alpha prior into one projected view.""" + if occupancy_probability is None: + return None + if occupancy_probability.shape != (view_batch, 1, UV_SIZE, UV_SIZE): + raise ValueError( + "outer_uv_occupancy probability must be grouped as " + f"{(view_batch, 1, UV_SIZE, UV_SIZE)}, got " + f"{tuple(occupancy_probability.shape)}." + ) + probability = occupancy_probability.float().flatten(2) + flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape( + 1, 1, -1 + ).expand(view_batch, 1, -1) + projected = probability.gather(2, flat_uv).reshape( + view_batch, *static["flat_uv"].shape[-2:] + ) + outer_valid = static["masks"][ROUTE_OUTER_PRIMARY].unsqueeze(0) + return torch.where(outer_valid, projected, torch.zeros_like(projected)) + + +def _outer_graph_component_labels(candidate, edge_index, max_steps=64): + """Label candidate outer nodes by bounded min-label propagation.""" + batch, node_count = candidate.shape + labels = torch.arange( + node_count, + device=candidate.device, + dtype=torch.long, + ).view(1, -1).expand(batch, -1).clone() + labels[~candidate] = node_count + source, target = edge_index + expanded_target = target.view(1, -1).expand(batch, -1) + connected_edge = candidate[:, source] & candidate[:, target] + for _ in range(int(max_steps)): + neighbour = torch.where( + connected_edge, + labels[:, source], + torch.full_like(labels[:, source], node_count), + ) + updated = labels.clone() + updated.scatter_reduce_( + 1, + expanded_target, + neighbour, + reduce="amin", + include_self=True, + ) + labels = updated + return labels + + +def outer_uv_topology_hysteresis( + probability, + seed_threshold=0.80, + grow_threshold=0.50, + min_component_size=2, + max_steps=64, +): + """Grow high-confidence outer seeds through topology-supported candidates.""" + if probability.dim() == 3: + probability = probability.unsqueeze(1) + if probability.dim() != 4 or probability.shape[1:] != ( + 1, + UV_SIZE, + UV_SIZE, + ): + raise ValueError("Outer occupancy probability must be Bx1x64x64.") + if not 0.0 <= float(grow_threshold) <= float(seed_threshold) <= 1.0: + raise ValueError( + "Require 0 <= grow_threshold <= seed_threshold <= 1." + ) + if min_component_size < 1: + raise ValueError("min_component_size must be positive.") + flat_indices, edge_index = build_outer_uv_graph() + flat_indices = flat_indices.to(probability.device) + edge_index = edge_index.to(probability.device) + nodes = probability.float().flatten(2)[:, 0].index_select( + 1, flat_indices + ) + candidate = nodes >= float(grow_threshold) + seeds = nodes >= float(seed_threshold) + source, target = edge_index + expanded_target = target.view(1, -1).expand(nodes.shape[0], -1) + active = seeds.clone() + connected_edge = candidate[:, source] & candidate[:, target] + for _ in range(int(max_steps)): + neighbour_active = torch.where( + connected_edge, + active[:, source], + torch.zeros_like(active[:, source]), + ) + propagated_active = torch.zeros_like(active, dtype=torch.uint8) + propagated_active.scatter_reduce_( + 1, + expanded_target, + neighbour_active.to(torch.uint8), + reduce="amax", + include_self=False, + ) + active = active | (candidate & propagated_active.bool()) + + labels = _outer_graph_component_labels( + candidate, + edge_index, + max_steps=max_steps, + ) + node_count = nodes.shape[1] + batch_offsets = ( + torch.arange(nodes.shape[0], device=nodes.device).view(-1, 1) + * (node_count + 1) + ) + component_index = labels + batch_offsets + component_sizes = nodes.new_zeros(nodes.shape[0] * (node_count + 1)) + component_sizes.scatter_add_( + 0, + component_index[candidate], + torch.ones_like(nodes[candidate]), + ) + node_component_size = component_sizes[component_index] + active = active & (node_component_size >= int(min_component_size)) + grown = active & ~seeds + # Connectivity determines eligibility, not confidence. Propagating a + # single seed's score across the whole component made a 0.50 neighbour + # behave like a 0.90 texel and amplified one false seed into a large outer + # patch. Accepted nodes therefore retain their own calibrated score. + structured_nodes = torch.where(active, nodes, torch.zeros_like(nodes)) + + def scatter_nodes(node_values, dtype=None): + atlas = node_values.new_zeros( + node_values.shape[0], + UV_SIZE * UV_SIZE, + dtype=dtype, + ) + atlas[:, flat_indices] = node_values.to(dtype=atlas.dtype) + return atlas.reshape(-1, 1, UV_SIZE, UV_SIZE) + + return { + "probability": scatter_nodes(structured_nodes), + "accepted": scatter_nodes(active, dtype=torch.bool), + "seed": scatter_nodes(seeds & active, dtype=torch.bool), + "grown": scatter_nodes(grown, dtype=torch.bool), + "rejected_candidate": scatter_nodes( + candidate & ~active, + dtype=torch.bool, + ), + } + + +def aggregate_direct_outer_values_by_view( + renderer, + views, + values, + center_power=2.0, +): + """Pool image-space values into direct outer UV texels, keeping views separate.""" + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required.") + if values.dim() != 4: + raise ValueError("values must be shaped NCHW.") + if values.shape[0] % len(views) != 0: + raise ValueError("The value batch must be divisible by the view count.") + if center_power < 0.0: + raise ValueError("center_power must be non-negative.") + + group_count = values.shape[0] // len(views) + channel_count = values.shape[1] + uv_count = UV_SIZE * UV_SIZE + pooled_views = [] + supported_views = [] + for view_index, view in enumerate(views): + selection = slice(view_index, values.shape[0], len(views)) + selected = values[selection].float() + static = build_static_surface_routing(renderer, view, values.device) + if static["masks"].shape[-2:] != values.shape[-2:]: + raise ValueError( + f"View {view!r} mapping shape " + f"{tuple(static['masks'].shape[-2:])} does not match " + f"value shape {tuple(values.shape[-2:])}." + ) + + valid = static["masks"][ROUTE_OUTER_PRIMARY].reshape(1, 1, -1) + flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape(1, 1, -1) + center_weight = static["texel_center_score"][ + ROUTE_OUTER_PRIMARY + ].float().clamp(0.0, 1.0).pow(float(center_power)) + weight = center_weight.reshape(1, 1, -1) * valid + weight = weight.expand(group_count, 1, -1) + indices = flat_uv.expand(group_count, channel_count, -1) + + sums = selected.new_zeros(group_count, channel_count, uv_count) + sums.scatter_add_( + 2, + indices, + selected.flatten(2) * weight, + ) + counts = selected.new_zeros(group_count, 1, uv_count) + counts.scatter_add_( + 2, + flat_uv.expand(group_count, 1, -1), + weight, + ) + pooled_views.append(sums / counts.clamp_min(1e-6)) + supported_views.append(counts[:, 0] > 1e-6) + + return ( + torch.stack(pooled_views, dim=1), + torch.stack(supported_views, dim=1), + ) + + +def attach_projected_outer_uv_occupancy( + model, + outputs, + renderer, + views, + observed_foreground=None, + center_power=2.0, +): + """Project dense parser features into UV space and run the topology head.""" + if not getattr(model, "predict_outer_uv_occupancy", False): + return outputs + if "outer_uv_features" not in outputs: + raise ValueError( + "Projected occupancy parser did not return outer_uv_features." + ) + views = parse_views(views) + feature_map = outputs["outer_uv_features"] + layer_logits = outputs["layer"] + foreground = torch.sigmoid(outputs["foreground"].float()) + if "affine" in outputs: + affine = outputs["affine"] + feature_map = canonicalize_tensor(feature_map, affine, mode="bilinear") + layer_logits = canonicalize_tensor( + layer_logits, affine, mode="bilinear" + ) + foreground = canonicalize_tensor( + foreground, affine, mode="bilinear" + ) + if observed_foreground is not None: + if observed_foreground.dim() == 3: + observed_foreground = observed_foreground.unsqueeze(1) + observed_foreground = canonicalize_tensor( + observed_foreground.float(), + affine, + mode="bilinear", + ) + elif observed_foreground is not None and observed_foreground.dim() == 3: + observed_foreground = observed_foreground.unsqueeze(1) + if observed_foreground is not None: + foreground = observed_foreground.to( + device=foreground.device, + dtype=foreground.dtype, + ).clamp(0.0, 1.0) + + outer_probability = torch.softmax( + layer_logits.float(), dim=1 + )[:, ROUTE_OUTER_PRIMARY : ROUTE_OUTER_PRIMARY + 1].detach() + values = torch.cat( + [ + feature_map.float(), + outer_probability, + foreground.float().detach(), + ], + dim=1, + ) + pooled, support = aggregate_direct_outer_values_by_view( + renderer, + views, + values, + center_power=center_power, + ) + support_weight = support.unsqueeze(2).to(dtype=pooled.dtype) + support_count = support_weight.sum(dim=1).clamp_min(1.0) + feature_channels = feature_map.shape[1] + pooled_mean = ( + (pooled * support_weight).sum(dim=1) / support_count + ) + mean_features = pooled_mean[:, :feature_channels] + mean_outer = pooled_mean[:, feature_channels : feature_channels + 1] + mean_foreground = pooled_mean[:, feature_channels + 1 : feature_channels + 2] + per_view_outer = pooled[:, :, feature_channels : feature_channels + 1] + max_outer = torch.where( + support_weight.bool(), + per_view_outer, + torch.zeros_like(per_view_outer), + ).amax(dim=1) + route_evidence = torch.cat([mean_outer, max_outer], dim=1) + route_evidence_dropout = float( + getattr(model, "outer_uv_route_evidence_dropout", 0.0) + ) + if route_evidence_dropout >= 1.0: + # The projected ownership classifier must remain independent from the + # already-unreliable per-pixel route prediction. Keep the two channels + # for checkpoint compatibility, but disable the shortcut in both + # training and inference. + route_evidence = torch.zeros_like(route_evidence) + elif model.training and route_evidence_dropout > 0.0: + # Drop the entire route-evidence pair per grouped skin. This prevents + # the occupancy head from learning the shortcut "copy p_outer" while + # leaving half of the batches available for useful route context. + keep_route_evidence = ( + torch.rand( + (route_evidence.shape[0],) + + (1,) * (route_evidence.ndim - 1), + device=route_evidence.device, + ) + >= route_evidence_dropout + ) + route_evidence = route_evidence * keep_route_evidence.to( + dtype=route_evidence.dtype + ) + support_fraction = ( + support_count / max(float(len(views)), 1.0) + ) + atlas_features = torch.cat( + [ + mean_features, + route_evidence, + mean_foreground, + support_fraction, + ], + dim=1, + ).reshape( + mean_features.shape[0], + feature_channels + 4, + UV_SIZE, + UV_SIZE, + ) + logits = model.predict_projected_outer_uv_occupancy( + atlas_features, + outputs["outer_uv_global_context"], + ) + outputs["outer_uv_occupancy_logits"] = logits + outputs["outer_uv_occupancy_support"] = ( + support.any(dim=1).reshape(-1, 1, UV_SIZE, UV_SIZE) + ) + return outputs + + +def _project_grouped_outer_texels_to_views( + renderer, + views, + texel_values, +): + """Project grouped outer-UV values back into every fixed input view.""" + views = parse_views(views) + if texel_values.dim() != 2 or texel_values.shape[1] != UV_SIZE * UV_SIZE: + raise ValueError( + "texel_values must be shaped " + f"(B, {UV_SIZE * UV_SIZE})." + ) + projected = [] + for view in views: + static = build_static_surface_routing( + renderer, view, texel_values.device + ) + flat_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY].reshape(1, -1) + gathered = texel_values.gather( + 1, flat_uv.expand(texel_values.shape[0], -1) + ).reshape(texel_values.shape[0], *static["masks"].shape[-2:]) + valid = static["masks"][ROUTE_OUTER_PRIMARY].unsqueeze(0) + projected.append(torch.where(valid, gathered, torch.zeros_like(gathered))) + return torch.stack(projected, dim=1).reshape( + texel_values.shape[0] * len(views), + *projected[0].shape[-2:], + ) + + +def cross_view_outer_visibility_evidence( + renderer, + views, + role_prob, + observed_foreground, + center_power=2.0, + positive_confidence=0.70, + positive_margin=0.20, + negative_confidence=0.70, + negative_margin=0.20, + background_max_coverage=0.25, + min_views=2, +): + """Find shared outer texels with high-confidence cross-view contradictions. + + A contradiction requires at least one view that strongly supports outer and + another view that either sees background at the projected outer texel or + strongly supports inner. Weak predictions, missing mappings, and texels + visible from only one view never form a veto. + """ + for name, value in ( + ("positive_confidence", positive_confidence), + ("positive_margin", positive_margin), + ("negative_confidence", negative_confidence), + ("negative_margin", negative_margin), + ("background_max_coverage", background_max_coverage), + ): + if not 0.0 <= float(value) <= 1.0: + raise ValueError(f"{name} must be in [0, 1].") + if min_views < 2: + raise ValueError("min_views must be at least 2.") + if observed_foreground.shape != ( + role_prob.shape[0], + *role_prob.shape[-2:], + ): + raise ValueError("observed_foreground must match the role probability map.") + + values = torch.cat( + [ + role_prob.float(), + observed_foreground.unsqueeze(1).float(), + ], + dim=1, + ) + pooled, support = aggregate_direct_outer_values_by_view( + renderer, + views, + values, + center_power=center_power, + ) + role_count = role_prob.shape[1] + pooled_role = pooled[:, :, :role_count] + foreground_coverage = pooled[:, :, role_count].clamp(0.0, 1.0) + outer_score = pooled_role[:, :, ROUTE_OUTER_PRIMARY] + outer_competitor = torch.maximum( + pooled_role[:, :, ROUTE_INNER_PRIMARY], + pooled_role[:, :, ROUTE_SECONDARY], + ) + outer_margin_ratio = ( + (outer_score - outer_competitor).clamp_min(0.0) + / outer_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + inner_score = pooled_role[:, :, ROUTE_INNER_PRIMARY] + inner_competitor = torch.maximum( + pooled_role[:, :, ROUTE_OUTER_PRIMARY], + pooled_role[:, :, ROUTE_SECONDARY], + ) + inner_margin_ratio = ( + (inner_score - inner_competitor).clamp_min(0.0) + / inner_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + + positive = ( + support + & (foreground_coverage > float(background_max_coverage)) + & (outer_score >= float(positive_confidence)) + & (outer_margin_ratio >= float(positive_margin)) + ) + background_negative = support & ( + foreground_coverage <= float(background_max_coverage) + ) + inner_negative = ( + support + & ~background_negative + & (inner_score >= float(negative_confidence)) + & (inner_margin_ratio >= float(negative_margin)) + ) + negative = background_negative | inner_negative + shared = support.sum(dim=1) >= int(min_views) + conflict = shared & positive.any(dim=1) & negative.any(dim=1) + + evidence_weight = support.to(dtype=outer_score.dtype) + prior = ( + (outer_score * foreground_coverage * evidence_weight).sum(dim=1) + / evidence_weight.sum(dim=1).clamp_min(1.0) + ).clamp(0.0, 1.0) + return { + "prior": _project_grouped_outer_texels_to_views( + renderer, views, prior + ), + "shared": _project_grouped_outer_texels_to_views( + renderer, views, shared.float() + ) > 0.5, + "conflict": _project_grouped_outer_texels_to_views( + renderer, views, conflict.float() + ) > 0.5, + "positive": _project_grouped_outer_texels_to_views( + renderer, views, positive.any(dim=1).float() + ) > 0.5, + "negative": _project_grouped_outer_texels_to_views( + renderer, views, negative.any(dim=1).float() + ) > 0.5, + "shared_texels": shared, + "conflict_texels": conflict, + } + + +def _soft_route_role_consensus( + role_prob, + static, + foreground_prob, + consensus_weight=0.60, + preserve_outer_confidence=0.80, + preserve_outer_margin=0.35, + outer_confidence=0.80, + outer_margin=0.35, + outer_uv_occupancy=None, + occupancy_blend_weight=0.0, + occupancy_gate_threshold=0.15, + occupancy_rescue_threshold=0.70, + occupancy_rescue_route_threshold=0.30, + cross_view_outer_prior=None, + cross_view_outer_valid=None, + cross_view_outer_veto=None, + cross_view_outer_weight=0.50, +): + """Denoise route roles without erasing strong local outer-layer evidence.""" + for name, value in ( + ("consensus_weight", consensus_weight), + ("preserve_outer_confidence", preserve_outer_confidence), + ("preserve_outer_margin", preserve_outer_margin), + ("outer_confidence", outer_confidence), + ("outer_margin", outer_margin), + ("occupancy_blend_weight", occupancy_blend_weight), + ("occupancy_gate_threshold", occupancy_gate_threshold), + ("occupancy_rescue_threshold", occupancy_rescue_threshold), + ( + "occupancy_rescue_route_threshold", + occupancy_rescue_route_threshold, + ), + ("cross_view_outer_weight", cross_view_outer_weight), + ): + if not 0.0 <= float(value) <= 1.0: + raise ValueError(f"{name} must be in [0, 1].") + + role_prob = role_prob.float() + raw_role = role_prob.argmax(dim=1) + raw_top = role_prob.topk(2, dim=1).values + raw_margin = raw_top[:, 0] - raw_top[:, 1] + raw_selected_score = role_prob.gather(1, raw_role.unsqueeze(1)).squeeze(1) + raw_margin_ratio = ( + raw_margin / raw_selected_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + + direct_scores = _aggregate_role_probabilities_by_direct_texel( + role_prob, + static, + foreground_prob=foreground_prob, + center_power=2.0, + ) + batch = role_prob.shape[0] + inner_valid = static["masks"][0].unsqueeze(0).expand(batch, -1, -1) + outer_valid = static["masks"][1].unsqueeze(0).expand_as(inner_valid) + has_direct_candidate = inner_valid | outer_valid + + texel_inner = direct_scores[:, 0, ROUTE_INNER_PRIMARY] + texel_outer = direct_scores[:, 1, ROUTE_OUTER_PRIMARY] + unavailable = role_prob.new_full(texel_inner.shape, -1.0) + texel_secondary = torch.maximum( + torch.where( + inner_valid, + direct_scores[:, 0, ROUTE_SECONDARY], + unavailable, + ), + torch.where( + outer_valid, + direct_scores[:, 1, ROUTE_SECONDARY], + unavailable, + ), + ) + + local_weight = 1.0 - float(consensus_weight) + fused_inner = torch.where( + inner_valid, + local_weight * role_prob[:, ROUTE_INNER_PRIMARY] + + float(consensus_weight) * texel_inner, + unavailable, + ) + fused_outer = torch.where( + outer_valid, + local_weight * role_prob[:, ROUTE_OUTER_PRIMARY] + + float(consensus_weight) * texel_outer, + unavailable, + ) + if cross_view_outer_prior is not None: + if ( + cross_view_outer_valid is None + or cross_view_outer_prior.shape != texel_outer.shape + or cross_view_outer_valid.shape != texel_outer.shape + ): + raise ValueError( + "Cross-view outer prior and validity must match the route map." + ) + cross_view_outer_prior = cross_view_outer_prior.float().clamp(0.0, 1.0) + cross_view_outer_valid = cross_view_outer_valid.bool() + fused_outer = torch.where( + outer_valid & cross_view_outer_valid, + (1.0 - float(cross_view_outer_weight)) * fused_outer + + float(cross_view_outer_weight) * cross_view_outer_prior, + fused_outer, + ) + if cross_view_outer_veto is None: + cross_view_outer_veto = torch.zeros_like(outer_valid) + elif cross_view_outer_veto.shape != texel_outer.shape: + raise ValueError("Cross-view outer veto must match the route map.") + else: + cross_view_outer_veto = cross_view_outer_veto.bool() + occupancy_available = outer_uv_occupancy is not None + if occupancy_available: + if outer_uv_occupancy.shape != texel_outer.shape: + raise ValueError( + "Projected outer occupancy must match the route map, got " + f"{tuple(outer_uv_occupancy.shape)} and " + f"{tuple(texel_outer.shape)}." + ) + occupancy_prior = outer_uv_occupancy.float().clamp(0.0, 1.0) + fused_outer = torch.where( + outer_valid, + (1.0 - float(occupancy_blend_weight)) * fused_outer + + float(occupancy_blend_weight) * occupancy_prior, + unavailable, + ) + occupancy_allows_outer = ( + occupancy_prior >= float(occupancy_gate_threshold) + ) + else: + occupancy_prior = torch.ones_like(texel_outer) + occupancy_allows_outer = torch.ones_like(outer_valid) + # Secondary/backface observations remain strict cell-level decisions. The + # soft blend is needed for the ambiguous inner/outer pair; retaining an + # isolated secondary pixel would reintroduce exactly the speckle that + # projected-texel consensus is meant to remove. + fused_secondary = torch.where( + has_direct_candidate, + texel_secondary.clamp_min(0.0), + unavailable, + ) + + raw_outer_score = role_prob[:, ROUTE_OUTER_PRIMARY] + raw_outer_competitor = torch.maximum( + role_prob[:, ROUTE_INNER_PRIMARY], + role_prob[:, ROUTE_SECONDARY], + ) + raw_outer_margin = (raw_outer_score - raw_outer_competitor).clamp_min(0.0) + raw_outer_margin_ratio = ( + raw_outer_margin / raw_outer_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + raw_preserve_outer = ( + outer_valid + & ~cross_view_outer_veto + & occupancy_allows_outer + & (raw_role == ROUTE_OUTER_PRIMARY) + & (raw_outer_score >= float(preserve_outer_confidence)) + & (raw_outer_margin_ratio >= float(preserve_outer_margin)) + ) + + fused_outer_competitor = torch.maximum(fused_inner, fused_secondary) + fused_outer_margin = (fused_outer - fused_outer_competitor).clamp_min(0.0) + fused_outer_margin_ratio = ( + fused_outer_margin / fused_outer.clamp_min(1e-8) + ).clamp(0.0, 1.0) + fused_outer_admitted = ( + outer_valid + & (fused_outer >= float(outer_confidence)) + & (fused_outer_margin_ratio >= float(outer_margin)) + ) + # Raw outer confidence is deliberately not an admission shortcut. A + # coherent inner texture (eyes, mouth, shirt graphics) can make the local + # route head extremely confident while still being semantically wrong. + # Every ordinary outer decision must therefore pass the fused texel gate. + # Occupancy evidence remains the only explicit bypass because it predicts + # atlas-level outer existence rather than repeating the local route score. + outer_admitted = fused_outer_admitted & ( + (raw_role != ROUTE_OUTER_PRIMARY) | raw_preserve_outer + ) + preserve_outer = raw_preserve_outer & fused_outer_admitted + outer_gate_rejected = ( + outer_valid + & ~cross_view_outer_veto + & occupancy_allows_outer + & (raw_role == ROUTE_OUTER_PRIMARY) + & ~outer_admitted + ) + occupancy_rescue = ( + outer_valid + & occupancy_allows_outer + & (occupancy_prior >= float(occupancy_rescue_threshold)) + & ( + raw_outer_score + >= float(occupancy_rescue_route_threshold) + ) + ) if occupancy_available else torch.zeros_like(outer_valid) + allow_outer = ( + outer_valid + & ~cross_view_outer_veto + & occupancy_allows_outer + & (outer_admitted | occupancy_rescue) + ) + + outer_candidate_score = torch.where( + occupancy_rescue, + torch.maximum( + fused_outer, + occupancy_prior, + ), + fused_outer, + ) + eligible_outer = torch.where(allow_outer, outer_candidate_score, unavailable) + eligible_secondary = torch.where( + (fused_secondary >= SECONDARY_TEXEL_MIN_PROBABILITY) + & has_direct_candidate, + fused_secondary, + unavailable, + ) + eligible_scores = torch.stack( + [fused_inner, eligible_outer, eligible_secondary], + dim=1, + ) + selected_score, selected_role = eligible_scores.max(dim=1) + selected_role = torch.where(has_direct_candidate, selected_role, raw_role) + selected_score = torch.where( + has_direct_candidate, + selected_score, + raw_selected_score, + ) + top_scores = eligible_scores.topk(2, dim=1).values + selected_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) + selected_margin_ratio = ( + selected_margin / selected_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + + # If fused outer still wins but does not clear the asymmetric admission + # gate, keep its outer UV identity while marking it untrusted. Downstream + # filtering can then leave the observation unknown (and preserve it as + # rejected context) instead of writing an uncertain outer color into the + # inner atlas. When fused inner actually wins, the ordinary argmax above + # performs the intended outer-to-inner correction. + deferred_outer = outer_gate_rejected & ( + fused_outer >= fused_outer_competitor + ) + selected_role = torch.where( + deferred_outer, + torch.full_like(selected_role, ROUTE_OUTER_PRIMARY), + selected_role, + ) + selected_score = torch.where(deferred_outer, fused_outer, selected_score) + selected_margin = torch.where( + deferred_outer, + fused_outer_margin, + selected_margin, + ) + selected_margin_ratio = torch.where( + deferred_outer, + fused_outer_margin_ratio, + selected_margin_ratio, + ) + + return { + "role": selected_role, + "score": selected_score, + "margin": selected_margin, + "margin_ratio": selected_margin_ratio, + "scores": torch.stack( + [fused_inner, fused_outer, fused_secondary], + dim=1, + ), + "has_direct_candidate": has_direct_candidate, + "preserved_outer": preserve_outer, + "outer_gate_rejected": outer_gate_rejected, + "occupancy_rescued_outer": ( + occupancy_rescue + & (raw_role != ROUTE_OUTER_PRIMARY) + & (selected_role == ROUTE_OUTER_PRIMARY) + ), + "occupancy_rejected_outer": ( + outer_valid + & (raw_role == ROUTE_OUTER_PRIMARY) + & ~occupancy_allows_outer + ), + "cross_view_outer_veto": cross_view_outer_veto, + "outer_uv_occupancy": ( + occupancy_prior if occupancy_available else None + ), + "changed": has_direct_candidate & (selected_role != raw_role), + "inner_to_outer": ( + has_direct_candidate + & (raw_role == ROUTE_INNER_PRIMARY) + & (selected_role == ROUTE_OUTER_PRIMARY) + ), + "outer_to_inner": ( + has_direct_candidate + & (raw_role == ROUTE_OUTER_PRIMARY) + & (selected_role == ROUTE_INNER_PRIMARY) + ), + } + + +def _routing_from_geometry_outputs( + renderer, + views, + outputs, + fg_threshold=0.5, + outer_threshold=0.5, + texel_consensus=True, + texel_consensus_weight=0.60, + preserve_outer_confidence=0.80, + preserve_outer_margin=0.35, + consensus_outer_confidence=0.70, + consensus_outer_margin=0.20, + outer_uv_occupancy=False, + outer_uv_occupancy_blend_weight=0.0, + outer_uv_occupancy_gate_threshold=0.15, + outer_uv_occupancy_rescue_threshold=0.70, + outer_uv_occupancy_rescue_route_threshold=0.30, + outer_uv_component_routing=False, + outer_uv_component_seed_threshold=0.80, + outer_uv_component_grow_threshold=0.50, + outer_uv_component_min_size=2, + observed_foreground=None, + cross_view_outer_consistency=True, + cross_view_outer_weight=0.50, + cross_view_outer_positive_confidence=0.70, + cross_view_outer_positive_margin=0.20, + cross_view_outer_negative_confidence=0.70, + cross_view_outer_negative_margin=0.20, + cross_view_outer_background_max_coverage=0.25, + cross_view_outer_min_views=2, +): + """Route a fitted Steve render through fixed inner/outer cuboid UV maps.""" + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required for geometry routing.") + foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] + if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES: + raise ValueError( + "geometry_fit requires a three-class route-role head " + "(inner_primary, outer_primary, secondary_backface). Retrain the parser." + ) + role_prob = torch.softmax(outputs["layer"].float(), dim=1) + raw_route_role = role_prob.argmax(dim=1) + route_role = raw_route_role.clone() + outer_prob = role_prob[:, ROUTE_OUTER_PRIMARY].clone() + inner_prob = role_prob[:, ROUTE_INNER_PRIMARY].clone() + top_prob = role_prob.topk(2, dim=1).values + role_margin = top_prob[:, 0] - top_prob[:, 1] + N, H, W = outer_prob.shape + if N % len(views) != 0: + raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") + + fg = torch.zeros_like(outer_prob, dtype=torch.bool) + layer = torch.zeros((N, H, W), dtype=torch.long, device=outer_prob.device) + flat_uv = torch.zeros_like(layer) + surface = torch.full_like(layer, IGNORE_INDEX) + confidence = torch.zeros_like(outer_prob) + confidence_margin = torch.zeros_like(outer_prob) + confidence_margin_ratio = torch.zeros_like(outer_prob) + texel_center_score = torch.zeros_like(outer_prob) + part = torch.full_like(layer, IGNORE_INDEX) + face = torch.full_like(layer, IGNORE_INDEX) + consensus_changed = torch.zeros_like(fg) + consensus_inner_to_outer = torch.zeros_like(fg) + consensus_outer_to_inner = torch.zeros_like(fg) + consensus_preserved_outer = torch.zeros_like(fg) + consensus_outer_gate_rejected = torch.zeros_like(fg) + occupancy_rescued_outer = torch.zeros_like(fg) + occupancy_rejected_outer = torch.zeros_like(fg) + projected_outer_occupancy = torch.zeros_like(outer_prob) + projected_outer_component_seed = torch.zeros_like(fg) + projected_outer_component_grown = torch.zeros_like(fg) + projected_outer_component_rejected = torch.zeros_like(fg) + occupancy_logits = ( + outputs.get("outer_uv_occupancy_logits") + if outer_uv_occupancy + else None + ) + occupancy_details = None + occupancy_probability = None + if occupancy_logits is not None: + raw_occupancy_probability = torch.sigmoid(occupancy_logits.float()) + if outer_uv_component_routing: + occupancy_details = outer_uv_topology_hysteresis( + raw_occupancy_probability, + seed_threshold=outer_uv_component_seed_threshold, + grow_threshold=outer_uv_component_grow_threshold, + min_component_size=outer_uv_component_min_size, + ) + occupancy_probability = occupancy_details["probability"] + else: + occupancy_probability = raw_occupancy_probability + cross_view = None + if cross_view_outer_consistency and len(views) >= cross_view_outer_min_views: + if observed_foreground is None: + observed_foreground = foreground_prob > float(fg_threshold) + cross_view = cross_view_outer_visibility_evidence( + renderer, + views, + role_prob, + observed_foreground, + positive_confidence=cross_view_outer_positive_confidence, + positive_margin=cross_view_outer_positive_margin, + negative_confidence=cross_view_outer_negative_confidence, + negative_margin=cross_view_outer_negative_margin, + background_max_coverage=( + cross_view_outer_background_max_coverage + ), + min_views=cross_view_outer_min_views, + ) + for view_index, view in enumerate(views): + selection = slice(view_index, N, len(views)) + static = build_static_surface_routing(renderer, view, outer_prob.device) + if static["masks"].shape[-2:] != (H, W): + raise ValueError( + f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " + f"does not match parser shape {(H, W)}." + ) + + inner_valid = static["masks"][0].unsqueeze(0).expand(outer_prob[selection].shape[0], -1, -1) + outer_valid = static["masks"][1].unsqueeze(0).expand_as(inner_valid) + if ( + texel_consensus + or cross_view is not None + or occupancy_probability is not None + ): + view_occupancy = _project_outer_uv_occupancy_to_view( + occupancy_probability, + static, + outer_prob[selection].shape[0], + ) + if occupancy_details is not None: + projected_outer_component_seed[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["seed"].float(), + static, + outer_prob[selection].shape[0], + ) + > 0.5 + ) + projected_outer_component_grown[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["grown"].float(), + static, + outer_prob[selection].shape[0], + ) + > 0.5 + ) + projected_outer_component_rejected[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["rejected_candidate"].float(), + static, + outer_prob[selection].shape[0], + ) + > 0.5 + ) + consensus = _soft_route_role_consensus( + role_prob[selection], + static, + foreground_prob[selection], + consensus_weight=( + texel_consensus_weight if texel_consensus else 0.0 + ), + preserve_outer_confidence=preserve_outer_confidence, + preserve_outer_margin=preserve_outer_margin, + outer_confidence=consensus_outer_confidence, + outer_margin=consensus_outer_margin, + outer_uv_occupancy=view_occupancy, + occupancy_blend_weight=( + outer_uv_occupancy_blend_weight + ), + occupancy_gate_threshold=( + outer_uv_occupancy_gate_threshold + ), + occupancy_rescue_threshold=( + outer_uv_occupancy_rescue_threshold + ), + occupancy_rescue_route_threshold=( + outer_uv_occupancy_rescue_route_threshold + ), + cross_view_outer_prior=( + cross_view["prior"][selection] + if cross_view is not None + else None + ), + cross_view_outer_valid=( + cross_view["shared"][selection] + if cross_view is not None + else None + ), + cross_view_outer_veto=( + cross_view["conflict"][selection] + if cross_view is not None + else None + ), + cross_view_outer_weight=cross_view_outer_weight, + ) + has_direct_candidate = consensus["has_direct_candidate"] + route_role[selection] = consensus["role"] + inner_prob[selection] = torch.where( + has_direct_candidate, + torch.where( + consensus["role"] == ROUTE_INNER_PRIMARY, + consensus["score"], + consensus["scores"][:, ROUTE_INNER_PRIMARY], + ).clamp_min(0.0), + inner_prob[selection], + ) + outer_prob[selection] = torch.where( + has_direct_candidate, + torch.where( + consensus["role"] == ROUTE_OUTER_PRIMARY, + consensus["score"], + consensus["scores"][:, ROUTE_OUTER_PRIMARY], + ).clamp_min(0.0), + outer_prob[selection], + ) + role_margin[selection] = consensus["margin"] + consensus_changed[selection] = consensus["changed"] + consensus_inner_to_outer[selection] = consensus["inner_to_outer"] + consensus_outer_to_inner[selection] = consensus["outer_to_inner"] + consensus_preserved_outer[selection] = consensus["preserved_outer"] + consensus_outer_gate_rejected[selection] = consensus[ + "outer_gate_rejected" + ] + occupancy_rescued_outer[selection] = consensus[ + "occupancy_rescued_outer" + ] + occupancy_rejected_outer[selection] = consensus[ + "occupancy_rejected_outer" + ] + if view_occupancy is not None: + projected_outer_occupancy[selection] = view_occupancy + + visible_outer = ( + (route_role[selection] == ROUTE_OUTER_PRIMARY) + & (outer_prob[selection] >= outer_threshold) + ) + # The Steve base layer is opaque and therefore supplies a complete + # silhouette. Outer-only protrusions still need image foreground evidence. + choose_outer = outer_valid & visible_outer & ( + inner_valid | (foreground_prob[selection] > fg_threshold) + ) + choose_inner = inner_valid & (route_role[selection] == ROUTE_INNER_PRIMARY) + selected_fg = choose_inner | choose_outer + + inner_flat_uv = static["flat_uv"][0].unsqueeze(0).expand_as(layer[selection]) + outer_flat_uv = static["flat_uv"][1].unsqueeze(0).expand_as(layer[selection]) + inner_part = static["part"][0].unsqueeze(0).expand_as(layer[selection]) + outer_part = static["part"][1].unsqueeze(0).expand_as(layer[selection]) + inner_face = static["face"][0].unsqueeze(0).expand_as(layer[selection]) + outer_face = static["face"][1].unsqueeze(0).expand_as(layer[selection]) + inner_center_score = static["texel_center_score"][0].unsqueeze(0).expand_as( + confidence[selection] + ) + outer_center_score = static["texel_center_score"][1].unsqueeze(0).expand_as( + confidence[selection] + ) + + fg[selection] = selected_fg + layer[selection] = choose_outer.long() + flat_uv[selection] = torch.where(choose_outer, outer_flat_uv, inner_flat_uv) + surface[selection] = torch.where( + selected_fg, + choose_outer.long(), + torch.full_like(layer[selection], IGNORE_INDEX), + ) + part[selection] = torch.where( + selected_fg, + torch.where(choose_outer, outer_part, inner_part), + torch.full_like(layer[selection], IGNORE_INDEX), + ) + face[selection] = torch.where( + selected_fg, + torch.where(choose_outer, outer_face, inner_face), + torch.full_like(layer[selection], IGNORE_INDEX), + ) + selected_confidence = torch.where(choose_outer, outer_prob[selection], inner_prob[selection]) + margin = role_margin[selection] + confidence[selection] = selected_confidence + confidence_margin[selection] = margin + confidence_margin_ratio[selection] = ( + margin / selected_confidence.clamp_min(1e-8) + ).clamp(0.0, 1.0) + texel_center_score[selection] = torch.where( + choose_outer, outer_center_score, inner_center_score + ) + + secondary = (route_role == ROUTE_SECONDARY) & (foreground_prob > fg_threshold) + return { + "foreground": fg, + "layer": layer, + "flat_uv": flat_uv, + "surface": surface, + "confidence": confidence, + "confidence_margin": confidence_margin, + "confidence_margin_ratio": confidence_margin_ratio, + "semantic_fallback": torch.zeros_like(fg), + "part": part, + "face": face, + "raw_route_role": raw_route_role, + "route_role": route_role, + "secondary": secondary, + "texel_center_score": texel_center_score, + "consensus_changed": consensus_changed, + "consensus_inner_to_outer": consensus_inner_to_outer, + "consensus_outer_to_inner": consensus_outer_to_inner, + "consensus_preserved_outer": consensus_preserved_outer, + "consensus_outer_gate_rejected": consensus_outer_gate_rejected, + "occupancy_rescued_outer": occupancy_rescued_outer, + "occupancy_rejected_outer": occupancy_rejected_outer, + "projected_outer_occupancy": projected_outer_occupancy, + "projected_outer_component_seed": projected_outer_component_seed, + "projected_outer_component_grown": projected_outer_component_grown, + "projected_outer_component_rejected": ( + projected_outer_component_rejected + ), + "outer_uv_occupancy_available": torch.full_like( + fg, occupancy_logits is not None + ), + "cross_view_outer_shared": ( + cross_view["shared"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_conflict": ( + cross_view["conflict"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_vetoed": ( + cross_view["conflict"] + & (raw_route_role == ROUTE_OUTER_PRIMARY) + & (route_role != ROUTE_OUTER_PRIMARY) + if cross_view is not None + else torch.zeros_like(fg) + ), + "cross_view_outer_positive": ( + cross_view["positive"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_negative": ( + cross_view["negative"] if cross_view is not None else torch.zeros_like(fg) + ), + } + + +def _aggregate_surface_scores_by_candidate_texel(candidate_score, static): + """Average each surface candidate inside its projected Minecraft UV texel.""" + candidate_score = candidate_score.float() + batch, surface_count, height, width = candidate_score.shape + flat_uv = static["flat_uv"][:surface_count].reshape(1, surface_count, -1) + valid = static["masks"][:surface_count].reshape(1, surface_count, -1) + indices = flat_uv.expand(batch, -1, -1) + values = candidate_score.flatten(2) * valid.to(dtype=candidate_score.dtype) + + sums = candidate_score.new_zeros(batch, surface_count, UV_SIZE * UV_SIZE) + sums.scatter_add_(2, indices, values) + counts = candidate_score.new_zeros(1, surface_count, UV_SIZE * UV_SIZE) + counts.scatter_add_( + 2, + flat_uv, + valid.to(dtype=candidate_score.dtype), + ) + means = sums / counts.clamp_min(1.0) + return means.gather(2, indices).reshape(batch, surface_count, height, width) + + +def _static_candidate_roles(static): + """Return the Minecraft layer and route role implemented by every surface pixel.""" + candidate_layer = static["layer"].clamp(0, LAYER_CLASSES - 1) + direct_masks = static["masks"][:LAYER_CLASSES].gather(0, candidate_layer) + direct_uv = static["flat_uv"][:LAYER_CLASSES].gather(0, candidate_layer) + primary = static["masks"] & direct_masks & (static["flat_uv"] == direct_uv) + candidate_role = torch.where( + primary, + candidate_layer, + torch.full_like(candidate_layer, ROUTE_SECONDARY), + ) + return candidate_layer, candidate_role + + +def _routing_from_geometry_surface_outputs( + renderer, + views, + outputs, + fg_threshold=0.5, + texel_consensus=True, + texel_consensus_weight=0.60, + preserve_outer_confidence=0.80, + preserve_outer_margin=0.35, + consensus_outer_confidence=0.70, + consensus_outer_margin=0.20, + outer_uv_occupancy=False, + outer_uv_occupancy_blend_weight=0.0, + outer_uv_occupancy_gate_threshold=0.15, + outer_uv_occupancy_rescue_threshold=0.70, + outer_uv_occupancy_rescue_route_threshold=0.30, + outer_uv_component_routing=False, + outer_uv_component_seed_threshold=0.80, + outer_uv_component_grow_threshold=0.50, + outer_uv_component_min_size=2, + observed_foreground=None, + cross_view_outer_consistency=True, + cross_view_outer_weight=0.50, + cross_view_outer_positive_confidence=0.70, + cross_view_outer_positive_margin=0.20, + cross_view_outer_negative_confidence=0.70, + cross_view_outer_negative_margin=0.20, + cross_view_outer_background_max_coverage=0.25, + cross_view_outer_min_views=2, +): + """Route primary layers first, then use the surface head for secondary UVs. + + The route-role head owns the inner/outer/secondary decision. The surface + head only disambiguates renderer slots that implement the selected role. + Keeping those decisions hierarchical is important: many composite and + geometry slots overlap a direct cuboid face, so a strong wrong surface + logit must not turn a primary pixel into a secondary/back-facing pixel. + """ + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required for geometry-surface routing.") + if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES: + raise ValueError("Geometry-surface routing requires the three-class route-role head.") + if "surface" not in outputs: + raise ValueError("Geometry-surface routing requires a static surface classifier.") + + foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] + role_prob = torch.softmax(outputs["layer"].float(), dim=1) + surface_prob = torch.softmax(outputs["surface"].float(), dim=1) + raw_route_role = role_prob.argmax(dim=1) + N, _, H, W = surface_prob.shape + if N % len(views) != 0: + raise ValueError(f"N={N} must be divisible by {len(views)} renderer views.") + + fg = torch.zeros((N, H, W), dtype=torch.bool, device=surface_prob.device) + layer = torch.zeros((N, H, W), dtype=torch.long, device=surface_prob.device) + flat_uv = torch.zeros_like(layer) + surface = torch.full_like(layer, IGNORE_INDEX) + route_role = torch.full_like(layer, IGNORE_INDEX) + confidence = torch.zeros_like(foreground_prob) + confidence_margin = torch.zeros_like(foreground_prob) + confidence_margin_ratio = torch.zeros_like(foreground_prob) + texel_center_score = torch.zeros_like(foreground_prob) + part = torch.full_like(layer, IGNORE_INDEX) + face = torch.full_like(layer, IGNORE_INDEX) + consensus_changed = torch.zeros_like(fg) + consensus_inner_to_outer = torch.zeros_like(fg) + consensus_outer_to_inner = torch.zeros_like(fg) + consensus_preserved_outer = torch.zeros_like(fg) + consensus_outer_gate_rejected = torch.zeros_like(fg) + occupancy_rescued_outer = torch.zeros_like(fg) + occupancy_rejected_outer = torch.zeros_like(fg) + projected_outer_occupancy = torch.zeros_like(foreground_prob) + projected_outer_component_seed = torch.zeros_like(fg) + projected_outer_component_grown = torch.zeros_like(fg) + projected_outer_component_rejected = torch.zeros_like(fg) + occupancy_logits = ( + outputs.get("outer_uv_occupancy_logits") + if outer_uv_occupancy + else None + ) + occupancy_details = None + occupancy_probability = None + if occupancy_logits is not None: + raw_occupancy_probability = torch.sigmoid(occupancy_logits.float()) + if outer_uv_component_routing: + occupancy_details = outer_uv_topology_hysteresis( + raw_occupancy_probability, + seed_threshold=outer_uv_component_seed_threshold, + grow_threshold=outer_uv_component_grow_threshold, + min_component_size=outer_uv_component_min_size, + ) + occupancy_probability = occupancy_details["probability"] + else: + occupancy_probability = raw_occupancy_probability + cross_view = None + if cross_view_outer_consistency and len(views) >= cross_view_outer_min_views: + if observed_foreground is None: + observed_foreground = foreground_prob > float(fg_threshold) + cross_view = cross_view_outer_visibility_evidence( + renderer, + views, + role_prob, + observed_foreground, + positive_confidence=cross_view_outer_positive_confidence, + positive_margin=cross_view_outer_positive_margin, + negative_confidence=cross_view_outer_negative_confidence, + negative_margin=cross_view_outer_negative_margin, + background_max_coverage=( + cross_view_outer_background_max_coverage + ), + min_views=cross_view_outer_min_views, + ) + + for view_index, view in enumerate(views): + selection = slice(view_index, N, len(views)) + static = build_static_surface_routing(renderer, view, surface_prob.device) + if static["masks"].shape[-2:] != (H, W): + raise ValueError( + f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " + f"does not match parser shape {(H, W)}." + ) + surface_count = static["masks"].shape[0] + if surface_prob.shape[1] < surface_count: + raise ValueError( + f"Parser has {surface_prob.shape[1]} surface classes, but view {view!r} " + f"requires {surface_count}." + ) + + candidate_layer, candidate_role = _static_candidate_roles(static) + + view_batch = surface_prob[selection].shape[0] + view_role_prob = role_prob[selection] + selected_role = raw_route_role[selection] + selected_role_score = view_role_prob.gather( + 1, selected_role.unsqueeze(1) + ).squeeze(1) + role_top = view_role_prob.topk(2, dim=1).values + role_margin = role_top[:, 0] - role_top[:, 1] + + if ( + texel_consensus + or cross_view is not None + or occupancy_probability is not None + ): + view_occupancy = _project_outer_uv_occupancy_to_view( + occupancy_probability, + static, + view_batch, + ) + if occupancy_details is not None: + projected_outer_component_seed[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["seed"].float(), + static, + view_batch, + ) + > 0.5 + ) + projected_outer_component_grown[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["grown"].float(), + static, + view_batch, + ) + > 0.5 + ) + projected_outer_component_rejected[selection] = ( + _project_outer_uv_occupancy_to_view( + occupancy_details["rejected_candidate"].float(), + static, + view_batch, + ) + > 0.5 + ) + consensus = _soft_route_role_consensus( + view_role_prob, + static, + foreground_prob[selection], + consensus_weight=( + texel_consensus_weight if texel_consensus else 0.0 + ), + preserve_outer_confidence=preserve_outer_confidence, + preserve_outer_margin=preserve_outer_margin, + outer_confidence=consensus_outer_confidence, + outer_margin=consensus_outer_margin, + outer_uv_occupancy=view_occupancy, + occupancy_blend_weight=( + outer_uv_occupancy_blend_weight + ), + occupancy_gate_threshold=( + outer_uv_occupancy_gate_threshold + ), + occupancy_rescue_threshold=( + outer_uv_occupancy_rescue_threshold + ), + occupancy_rescue_route_threshold=( + outer_uv_occupancy_rescue_route_threshold + ), + cross_view_outer_prior=( + cross_view["prior"][selection] + if cross_view is not None + else None + ), + cross_view_outer_valid=( + cross_view["shared"][selection] + if cross_view is not None + else None + ), + cross_view_outer_veto=( + cross_view["conflict"][selection] + if cross_view is not None + else None + ), + cross_view_outer_weight=cross_view_outer_weight, + ) + selected_role = consensus["role"] + selected_role_score = consensus["score"] + role_margin = consensus["margin"] + consensus_changed[selection] = consensus["changed"] + consensus_inner_to_outer[selection] = consensus["inner_to_outer"] + consensus_outer_to_inner[selection] = consensus["outer_to_inner"] + consensus_preserved_outer[selection] = consensus["preserved_outer"] + consensus_outer_gate_rejected[selection] = consensus[ + "outer_gate_rejected" + ] + occupancy_rescued_outer[selection] = consensus[ + "occupancy_rescued_outer" + ] + occupancy_rejected_outer[selection] = consensus[ + "occupancy_rejected_outer" + ] + if view_occupancy is not None: + projected_outer_occupancy[selection] = view_occupancy + + candidate_score = surface_prob[selection, :surface_count] + candidate_valid = static["masks"].unsqueeze(0).expand( + view_batch, -1, -1, -1 + ) + candidate_valid = candidate_valid & ( + candidate_role.unsqueeze(0) == selected_role.unsqueeze(1) + ) + if texel_consensus: + consensus_score = _aggregate_surface_scores_by_candidate_texel( + candidate_score, + static, + ) + local_scores = candidate_score.masked_fill(~candidate_valid, -1.0) + local_top = local_scores.topk(k=min(2, surface_count), dim=1).values + if surface_count > 1: + local_margin_ratio = ( + (local_top[:, 0] - local_top[:, 1]).clamp_min(0.0) + / local_top[:, 0].abs().clamp_min(1e-8) + ).clamp(0.0, 1.0) + consensus_weight = (1.0 - local_margin_ratio).unsqueeze(1) + candidate_score = ( + candidate_score * (1.0 - consensus_weight) + + consensus_score * consensus_weight + ) + candidate_score = candidate_score.masked_fill(~candidate_valid, -1.0) + best_surface_score, selected_surface = candidate_score.max(dim=1) + selected_secondary = selected_role == ROUTE_SECONDARY + # Primary routes already have an exact direct cuboid surface and UV. + # Do not let an equivalent composite slot bypass direct-layer coverage + # filtering; the learned surface head is needed only for secondary UVs. + selected_surface = torch.where( + selected_secondary, + selected_surface, + selected_role, + ) + if surface_count > 1: + top_scores = candidate_score.topk(k=2, dim=1).values + surface_margin = (top_scores[:, 0] - top_scores[:, 1]).clamp_min(0.0) + surface_margin_ratio = ( + surface_margin / top_scores[:, 0].abs().clamp_min(1e-8) + ).clamp(0.0, 1.0) + else: + surface_margin = best_surface_score.clamp_min(0.0) + surface_margin_ratio = (best_surface_score > 0.0).to( + best_surface_score.dtype + ) + + routed = _select_static_surface(static, selected_surface) + has_candidate = best_surface_score >= 0.0 + base_silhouette = static["masks"][0].unsqueeze(0).expand(view_batch, -1, -1) + routed_fg = ( + ((foreground_prob[selection] > fg_threshold) | base_silhouette) + & has_candidate + & routed["valid"] + ) + role_margin_ratio = ( + role_margin / selected_role_score.clamp_min(1e-8) + ).clamp(0.0, 1.0) + selected_surface_score = best_surface_score.clamp_min(0.0) + selected_confidence = torch.where( + selected_secondary, + (selected_role_score * selected_surface_score).sqrt(), + selected_role_score, + ) + if "route_confidence" in outputs: + # The learned trust target already says whether the argmax route is + # correct (and includes surface correctness for secondary routes). + # Multiplying it by the class probability counts the same + # uncertainty twice and makes a 0.80 outer gate nearly + # unreachable. A geometric mean keeps both signals necessary + # without systematically collapsing otherwise well-calibrated + # confidence. + learned_trust = torch.sigmoid( + outputs["route_confidence"][selection, 0].float() + ) + primary_confidence = ( + selected_role_score * learned_trust + ).clamp_min(0.0).sqrt() + secondary_confidence = ( + selected_role_score + * selected_surface_score + * learned_trust + ).clamp_min(0.0).pow(1.0 / 3.0) + selected_confidence = torch.where( + selected_secondary, + secondary_confidence, + primary_confidence, + ) + selected_margin = torch.where( + selected_secondary, + torch.minimum(role_margin, surface_margin), + role_margin, + ) + selected_margin_ratio = torch.where( + selected_secondary, + torch.minimum(role_margin_ratio, surface_margin_ratio), + role_margin_ratio, + ) + + fg[selection] = routed_fg + layer[selection] = routed["layer"] + flat_uv[selection] = routed["flat_uv"] + surface[selection] = torch.where( + routed_fg, + selected_surface, + torch.full_like(selected_surface, IGNORE_INDEX), + ) + route_role[selection] = torch.where( + routed_fg, + selected_role, + torch.full_like(selected_role, IGNORE_INDEX), + ) + confidence[selection] = selected_confidence + confidence_margin[selection] = selected_margin + confidence_margin_ratio[selection] = selected_margin_ratio + part[selection] = routed["part"] + face[selection] = routed["face"] + texel_center_score[selection] = routed["texel_center_score"] + + secondary = fg & (route_role == ROUTE_SECONDARY) + return { + "foreground": fg, + "layer": layer, + "flat_uv": flat_uv, + "surface": surface, + "confidence": confidence, + "confidence_margin": confidence_margin, + "confidence_margin_ratio": confidence_margin_ratio, + "semantic_fallback": torch.zeros_like(fg), + "part": part, + "face": face, + "raw_route_role": raw_route_role, + "route_role": route_role, + "secondary": secondary, + "secondary_routed": secondary, + "texel_center_score": texel_center_score, + "consensus_changed": consensus_changed, + "consensus_inner_to_outer": consensus_inner_to_outer, + "consensus_outer_to_inner": consensus_outer_to_inner, + "consensus_preserved_outer": consensus_preserved_outer, + "consensus_outer_gate_rejected": consensus_outer_gate_rejected, + "occupancy_rescued_outer": occupancy_rescued_outer, + "occupancy_rejected_outer": occupancy_rejected_outer, + "projected_outer_occupancy": projected_outer_occupancy, + "projected_outer_component_seed": projected_outer_component_seed, + "projected_outer_component_grown": projected_outer_component_grown, + "projected_outer_component_rejected": ( + projected_outer_component_rejected + ), + "outer_uv_occupancy_available": torch.full_like( + fg, occupancy_logits is not None + ), + "cross_view_outer_shared": ( + cross_view["shared"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_conflict": ( + cross_view["conflict"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_vetoed": ( + cross_view["conflict"] + & (raw_route_role == ROUTE_OUTER_PRIMARY) + & (route_role != ROUTE_OUTER_PRIMARY) + if cross_view is not None + else torch.zeros_like(fg) + ), + "cross_view_outer_positive": ( + cross_view["positive"] if cross_view is not None else torch.zeros_like(fg) + ), + "cross_view_outer_negative": ( + cross_view["negative"] if cross_view is not None else torch.zeros_like(fg) + ), + "learned_trust": ( + torch.sigmoid(outputs["route_confidence"][:, 0].float()) + if "route_confidence" in outputs + else torch.ones_like(foreground_prob) + ), + } + + +def _surface_aware_outer_coverage(routing, trusted, selected_outer, renderer, views): + """Measure outer support within the selected view/surface/UV texel.""" + views = parse_views(views) + views_per_group = len(views) + group_count = trusted.shape[0] // views_per_group + static_views = [ + build_static_surface_routing(renderer, view, trusted.device) + for view in views + ] + surface_count = max(static["masks"].shape[0] for static in static_views) + uv_count = UV_SIZE * UV_SIZE + view_stride = surface_count * uv_count + group_stride = views_per_group * view_stride + expected = routing["confidence"].new_zeros(group_stride) + + for view_index, static in enumerate(static_views): + static_surface_count = static["masks"].shape[0] + surface_offsets = ( + torch.arange(static_surface_count, device=trusted.device).view(-1, 1, 1) + * uv_count + ) + expected_indices = ( + view_index * view_stride + + surface_offsets + + static["flat_uv"][:static_surface_count] + ) + expected_mask = ( + static["masks"][:static_surface_count] + & (static["layer"][:static_surface_count] == 1) + ) + selected_indices = expected_indices[expected_mask] + expected.scatter_add_( + 0, + selected_indices, + torch.ones( + selected_indices.shape[0], + device=trusted.device, + dtype=expected.dtype, + ), + ) + + item_indices = torch.arange(trusted.shape[0], device=trusted.device) + item_groups = item_indices // views_per_group + item_views = item_indices % views_per_group + selected_surface = routing["surface"].clamp(0, surface_count - 1) + route_index = ( + item_views.view(-1, 1, 1) * view_stride + + selected_surface * uv_count + + routing["flat_uv"] + ) + grouped_route_index = route_index + item_groups.view(-1, 1, 1) * group_stride + observed_mask = trusted & selected_outer + observed = routing["confidence"].new_zeros(group_count * group_stride) + observed_indices = grouped_route_index[observed_mask] + observed.scatter_add_( + 0, + observed_indices, + torch.ones( + observed_indices.shape[0], + device=trusted.device, + dtype=observed.dtype, + ), + ) + coverage = observed.reshape(group_count, group_stride) / expected.clamp_min(1.0) + pixel_coverage = coverage[item_groups.view(-1, 1, 1), route_index] + # Composite and geometry slots can be partially occluded by nearer surfaces, + # so their static masks are not valid coverage denominators. Their exact slot + # classifier and texel consensus already reject isolated fragments. + return torch.where( + selected_surface >= LAYER_CLASSES, + torch.ones_like(pixel_coverage), + pixel_coverage, + ) + + +def _outer_silhouette_coverage( + observed_foreground, + flat_uv, + renderer, + views, + dilation=0, + min_pixels=4, +): + """Measure whether an outer texel's protruding pixels exist in the cutout. + + Only pixels covered by the outer cuboid but not by the inner cuboid carry + useful silhouette evidence. Interior projections are deliberately left + unassessed because an absent outer texel and the inner surface behind it + have the same foreground mask there. + + Coverage is calculated independently for every view and the least + supported assessed view wins. Averaging views would let a texel that is + visible in one view but absent in another pass at exactly 50% coverage, + which creates the characteristic head-top and clothing-edge spikes. + """ + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required.") + if observed_foreground.shape != flat_uv.shape: + raise ValueError( + "observed_foreground and flat_uv must have identical shapes." + ) + if observed_foreground.shape[0] % len(views) != 0: + raise ValueError("The routed batch must contain complete view groups.") + if dilation < 0: + raise ValueError("dilation must be non-negative.") + if min_pixels < 1: + raise ValueError("min_pixels must be positive.") + + observed = observed_foreground.to(dtype=torch.bool) + if dilation > 0: + kernel_size = 2 * int(dilation) + 1 + observed = F.max_pool2d( + observed.unsqueeze(1).float(), + kernel_size=kernel_size, + stride=1, + padding=int(dilation), + )[:, 0] > 0.5 + + views_per_group = len(views) + group_count = observed.shape[0] // views_per_group + uv_count = UV_SIZE * UV_SIZE + expected = observed.new_zeros( + (views_per_group, uv_count), dtype=torch.float32 + ) + supported = observed.new_zeros( + (group_count, views_per_group, uv_count), dtype=torch.float32 + ) + for view_index, view in enumerate(views): + static = build_static_surface_routing( + renderer, view, observed.device + ) + inner_silhouette = static["masks"][ROUTE_INNER_PRIMARY] + outer_silhouette = static["masks"][ROUTE_OUTER_PRIMARY] + protruding = outer_silhouette & ~inner_silhouette + if not protruding.any(): + continue + protruding_uv = static["flat_uv"][ROUTE_OUTER_PRIMARY][ + protruding + ] + expected[view_index].scatter_add_( + 0, + protruding_uv, + torch.ones_like(protruding_uv, dtype=expected.dtype), + ) + observed_view = observed[view_index::views_per_group, protruding] + supported[:, view_index].scatter_add_( + 1, + protruding_uv.view(1, -1).expand(group_count, -1), + observed_view.to(dtype=supported.dtype), + ) + + assessed_by_view = expected >= float(min_pixels) + coverage_by_view = supported / expected.clamp_min(1.0).unsqueeze(0) + coverage_by_view = torch.where( + assessed_by_view.unsqueeze(0), + coverage_by_view, + torch.ones_like(coverage_by_view), + ) + coverage_uv = coverage_by_view.amin(dim=1) + assessed_uv = assessed_by_view.any(dim=0) + item_groups = ( + torch.arange(observed.shape[0], device=observed.device) + // views_per_group + ) + pixel_coverage = coverage_uv[ + item_groups.view(-1, 1, 1), flat_uv + ] + pixel_assessed = assessed_uv[flat_uv] + return pixel_coverage, pixel_assessed + + +def _grouped_uv_support(mask, flat_uv, group_size): + """Broadcast per-pixel evidence to matching UV texels across all input views.""" + if mask.shape != flat_uv.shape: + raise ValueError("mask and flat_uv must have identical shapes.") + if group_size <= 0 or mask.shape[0] % group_size != 0: + raise ValueError("group_size must evenly divide the routed batch.") + + group_count = mask.shape[0] // group_size + item_groups = torch.arange(mask.shape[0], device=mask.device) // group_size + grouped_uv = flat_uv + item_groups.view(-1, 1, 1) * (UV_SIZE * UV_SIZE) + support_count = torch.zeros( + group_count * UV_SIZE * UV_SIZE, + device=mask.device, + dtype=torch.float32, + ) + selected_uv = grouped_uv[mask] + support_count.scatter_add_( + 0, + selected_uv, + torch.ones(selected_uv.shape[0], device=mask.device, dtype=support_count.dtype), + ) + return support_count[grouped_uv] > 0 + + +def _geometry_supported_outer_texels(routing, renderer, views): + """Find outer texels proven by outer-only silhouette or an exact secondary slot. + + A direct outer pixel outside the base cuboid silhouette cannot be explained by + the inner layer. Likewise a routed secondary surface is tied to an exact + renderer slot. Either observation is safe evidence for relaxing confidence + gates on the same outer UV texel in another pixel or view. + """ + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required for outer rescue.") + selected_outer = routing["layer"] == 1 + if selected_outer.shape[0] % len(views) != 0: + raise ValueError("Routed batch must be divisible by the number of views.") + + outer_only = torch.zeros_like(selected_outer) + selected_surface = routing.get("surface") + for view_index, view in enumerate(views): + selection = slice(view_index, selected_outer.shape[0], len(views)) + static = build_static_surface_routing(renderer, view, selected_outer.device) + direct_outer = static["masks"][1] & ~static["masks"][0] + if selected_surface is not None: + direct_outer = direct_outer.unsqueeze(0) & ( + selected_surface[selection] == ROUTE_OUTER_PRIMARY + ) + else: + direct_outer = direct_outer.unsqueeze(0).expand( + selected_outer[selection].shape[0], -1, -1 + ) + outer_only[selection] = direct_outer + + secondary = routing.get("secondary", torch.zeros_like(selected_outer)) + geometry_seed = selected_outer & (outer_only | secondary) + supported_texel = _grouped_uv_support( + geometry_seed, + routing["flat_uv"], + len(views), + ) + return selected_outer & supported_texel + + +def _semantic_supported_outer_pixels( + outputs, + routing, + group_size, + presence_threshold=0.80, + coverage_threshold=0.20, +): + """Use image-level part semantics to identify plausible outer-layer pixels.""" + selected_outer = routing["layer"] == 1 + presence_logits = outputs.get("outer_presence_logits") + predicted_coverage = outputs.get("outer_coverage") + routed_part = routing.get("part") + if ( + presence_logits is None + or predicted_coverage is None + or routed_part is None + ): + return torch.zeros_like(selected_outer) + group_count = selected_outer.shape[0] // group_size + if presence_logits.shape[:2] != (group_count, PART_CLASSES): + return torch.zeros_like(selected_outer) + if predicted_coverage.shape[:2] != (group_count, PART_CLASSES): + return torch.zeros_like(selected_outer) + + item_groups = torch.arange( + selected_outer.shape[0], device=selected_outer.device + ) // group_size + safe_part = routed_part.clamp(0, PART_CLASSES - 1) + valid_part = routed_part != IGNORE_INDEX + expanded_groups = item_groups.view(-1, 1, 1).expand_as(safe_part) + presence = torch.sigmoid(presence_logits.float())[ + expanded_groups, safe_part + ] + coverage = predicted_coverage.float()[expanded_groups, safe_part] + return ( + selected_outer + & valid_part + & (presence >= float(presence_threshold)) + & (coverage >= float(coverage_threshold)) + ) + + +def _scatter_soft_uv(rgb_sum, weight_sum, source_rgb, weight, flat_uv, valid): + """Differentiably accumulate weighted render pixels into fixed UV indices.""" + if not valid.any(): + return rgb_sum, weight_sum + batch = source_rgb.shape[0] + indices = flat_uv[valid].reshape(1, -1).expand(batch, -1) + selected_weight = weight[:, valid] + selected_rgb = source_rgb[:, :, valid] * selected_weight.unsqueeze(1) + weight_contribution = weight_sum.new_zeros(batch, UV_SIZE * UV_SIZE).scatter_add( + 1, + indices, + selected_weight, + ) + rgb_contribution = rgb_sum.new_zeros(batch, 3, UV_SIZE * UV_SIZE).scatter_add( + 2, + indices.unsqueeze(1).expand(-1, 3, -1), + selected_rgb, + ) + return rgb_sum + rgb_contribution, weight_sum + weight_contribution + + +def soft_splat_geometry_predictions_to_uv( + rendered, + outputs, + renderer, + views, + group_size=None, + temperature=1.0, + canonicalize=True, + eps=1e-6, + return_details=False, +): + """Build a differentiable provisional UV atlas from geometry-parser probabilities. + + Primary inner/outer probabilities are splatted through their direct cuboid + maps. Secondary probability is distributed through the learned exact + surface probabilities. UV scatter indices stay fixed while all weights and + the predicted affine remain differentiable. + """ + if rendered.dim() != 4: + raise ValueError(f"Expected rendered NCHW tensor, got {tuple(rendered.shape)}.") + views = parse_views(views) + if not views: + raise ValueError("At least one renderer view is required for soft geometry splatting.") + if group_size is None: + group_size = len(views) + if group_size != len(views): + raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") + if rendered.shape[0] % group_size != 0: + raise ValueError(f"N={rendered.shape[0]} must be divisible by group_size={group_size}.") + if temperature <= 0: + raise ValueError("temperature must be positive.") + if outputs["layer"].shape[1] != ROUTE_ROLE_CLASSES or "surface" not in outputs: + raise ValueError( + "Soft geometry splatting requires three route-role logits and exact surface logits." + ) + + if canonicalize: + canonical_outputs = canonicalize_parser_outputs(outputs) + canonical_rendered = canonicalize_parser_render( + rendered, + outputs, + mode="bilinear", + ).float() + else: + canonical_outputs = outputs + canonical_rendered = rendered.float() + role_prob = torch.softmax(canonical_outputs["layer"].float() / temperature, dim=1) + surface_logits = canonical_outputs["surface"] + surface_log_normalizer = torch.logsumexp(surface_logits / temperature, dim=1) + foreground_prob = torch.sigmoid(canonical_outputs["foreground"][:, 0].float()) + groups = rendered.shape[0] // group_size + reference = canonical_rendered + rgb_sums = [ + reference.new_zeros(groups, 3, UV_SIZE * UV_SIZE) + for _ in range(LAYER_CLASSES) + ] + weight_sums = [ + reference.new_zeros(groups, UV_SIZE * UV_SIZE) + for _ in range(LAYER_CLASSES) + ] + expected = reference.new_zeros(LAYER_CLASSES, UV_SIZE * UV_SIZE) + + for view_index, view in enumerate(views): + selection = slice(view_index, rendered.shape[0], group_size) + static = build_static_surface_routing(renderer, view, rendered.device) + if static["masks"].shape[-2:] != canonical_rendered.shape[-2:]: + raise ValueError( + f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} " + f"does not match parser shape {tuple(canonical_rendered.shape[-2:])}." + ) + surface_count = static["masks"].shape[0] + if surface_logits.shape[1] < surface_count: + raise ValueError( + f"Parser has {surface_logits.shape[1]} surface classes, but view {view!r} " + f"requires {surface_count}." + ) + + source_rgb = canonical_rendered[selection, :3] + view_foreground = foreground_prob[selection] + view_role_prob = role_prob[selection] + view_surface_logits = surface_logits[selection] + view_surface_log_normalizer = surface_log_normalizer[selection] + + for layer_index in range(LAYER_CLASSES): + valid = static["masks"][layer_index] + flat_uv = static["flat_uv"][layer_index] + weight = view_foreground * view_role_prob[:, layer_index] + rgb_sums[layer_index], weight_sums[layer_index] = _scatter_soft_uv( + rgb_sums[layer_index], + weight_sums[layer_index], + source_rgb, + weight, + flat_uv, + valid, + ) + expected[layer_index].scatter_add_( + 0, + flat_uv[valid], + torch.ones_like(flat_uv[valid], dtype=expected.dtype), + ) + + candidate_layer, candidate_role = _static_candidate_roles(static) + for surface_index in range(LAYER_CLASSES, surface_count): + secondary = ( + static["masks"][surface_index] + & (candidate_role[surface_index] == ROUTE_SECONDARY) + ) + if not secondary.any(): + continue + surface_weight = ( + view_foreground + * view_role_prob[:, ROUTE_SECONDARY] + * torch.exp( + view_surface_logits[:, surface_index].float() / temperature + - view_surface_log_normalizer.float() + ) + ) + for layer_index in range(LAYER_CLASSES): + valid = secondary & (candidate_layer[surface_index] == layer_index) + if not valid.any(): + continue + rgb_sums[layer_index], weight_sums[layer_index] = _scatter_soft_uv( + rgb_sums[layer_index], + weight_sums[layer_index], + source_rgb, + surface_weight, + static["flat_uv"][surface_index], + valid, + ) + + stacked_rgb_sum = torch.stack(rgb_sums, dim=1) + stacked_weight_sum = torch.stack(weight_sums, dim=1).unsqueeze(2) + expected = expected.view(1, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE).clamp_min(1.0) + layer_alpha = (stacked_weight_sum / expected).clamp(0.0, 1.0) + total_rgb_sum = stacked_rgb_sum.sum(dim=1) + total_weight = stacked_weight_sum.sum(dim=1) + rgb = total_rgb_sum / total_weight.clamp_min(eps) + alpha = layer_alpha.amax(dim=1) + pred_uv = torch.cat([rgb, alpha], dim=1).reshape(groups, 4, UV_SIZE, UV_SIZE) + support = total_weight.reshape(groups, 1, UV_SIZE, UV_SIZE) + if return_details: + return pred_uv, { + "support": support, + "layer_weight": stacked_weight_sum.reshape( + groups, LAYER_CLASSES, UV_SIZE, UV_SIZE + ), + "layer_alpha": layer_alpha.reshape( + groups, LAYER_CLASSES, UV_SIZE, UV_SIZE + ), + "canonical_rendered": canonical_rendered, + "canonical_outputs": canonical_outputs, + } + return pred_uv + + +def render_direct_uv(skins, renderer, view): + """Differentiably render the direct inner/outer cuboids without deep composites.""" + if skins.dim() != 4 or skins.shape[1:] != (4, UV_SIZE, UV_SIZE): + raise ValueError(f"Expected skins shaped (N, 4, 64, 64), got {tuple(skins.shape)}.") + batch = skins.shape[0] + dtype = skins.dtype + inner_grid = getattr(renderer, f"{view}_inner_grid").to(dtype=dtype) + outer_grid = getattr(renderer, f"{view}_outer_grid").to(dtype=dtype) + inner_mask = getattr(renderer, f"{view}_inner_mask").to(dtype=dtype) + outer_mask = getattr(renderer, f"{view}_outer_mask").to(dtype=dtype) + inner = F.grid_sample( + skins, + inner_grid.unsqueeze(0).expand(batch, -1, -1, -1), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) * inner_mask.view(1, 1, *inner_mask.shape) + outer = F.grid_sample( + skins, + outer_grid.unsqueeze(0).expand(batch, -1, -1, -1), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) * outer_mask.view(1, 1, *outer_mask.shape) + bg = renderer.bg_color.to(device=skins.device, dtype=dtype).view(1, 3, 1, 1) + inner_rgb = inner[:, 3:4] * inner[:, :3] + (1.0 - inner[:, 3:4]) * bg + rgb = outer[:, 3:4] * outer[:, :3] + (1.0 - outer[:, 3:4]) * inner_rgb + alpha = outer[:, 3:4] + (1.0 - outer[:, 3:4]) * inner[:, 3:4] + return torch.cat([rgb, alpha], dim=1) + + +def render_outer_alpha_uv(skins, renderer, view): + """Differentiably project only the outer-layer alpha into one view.""" + if skins.dim() != 4 or skins.shape[1:] != (4, UV_SIZE, UV_SIZE): + raise ValueError( + f"Expected skins shaped (N, 4, 64, 64), got {tuple(skins.shape)}." + ) + batch = skins.shape[0] + dtype = skins.dtype + outer_grid = getattr(renderer, f"{view}_outer_grid").to(dtype=dtype) + outer_mask = getattr(renderer, f"{view}_outer_mask").to(dtype=dtype) + return F.grid_sample( + skins[:, 3:4], + outer_grid.unsqueeze(0).expand(batch, -1, -1, -1), + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) * outer_mask.view(1, 1, *outer_mask.shape) + + +def flat_uv_to_uv01(flat_uv, dtype): + return torch.stack( + [ + (flat_uv.remainder(UV_SIZE) / float(UV_SIZE - 1)).to(dtype), + (flat_uv.div(UV_SIZE, rounding_mode="floor") / float(UV_SIZE - 1)).to(dtype), + ], + dim=1, + ) + + +def splat_parser_predictions_to_uv_conditioning( + rendered, + outputs, + renderer=None, + views=None, + group_size=1, + fg_threshold=0.5, + bg_color=(128, 128, 128), + semantic_gate=True, + affine_refine=True, + affine_refine_translation_px=8.0, + affine_refine_scale=0.0, + route_confidence_threshold=0.0, + route_margin_threshold=0.0, + outer_route_confidence_threshold=None, + outer_route_margin_threshold=None, + outer_uv_min_coverage=0.5, + outer_uv_min_source_pixels=1, + outer_silhouette_consistency=True, + outer_silhouette_min_coverage=0.50, + outer_silhouette_dilation=0, + outer_silhouette_min_pixels=4, + outer_geometry_rescue=False, + outer_semantic_rescue=True, + outer_semantic_presence_threshold=0.80, + outer_semantic_coverage_threshold=0.20, + outer_rescue_confidence_threshold=0.60, + outer_rescue_margin_threshold=0.25, + outer_rescue_min_coverage=0.10, + color_aggregation="grid_mode", + geometry_outer_threshold=0.5, + geometry_route_texel_consensus=True, + geometry_route_texel_consensus_weight=0.60, + geometry_route_preserve_outer_confidence=0.80, + geometry_route_preserve_outer_margin=0.35, + geometry_route_consensus_outer_confidence=0.70, + geometry_route_consensus_outer_margin=0.20, + geometry_cross_view_outer_consistency=False, + geometry_cross_view_outer_weight=0.50, + geometry_cross_view_outer_positive_confidence=0.70, + geometry_cross_view_outer_positive_margin=0.20, + geometry_cross_view_outer_negative_confidence=0.70, + geometry_cross_view_outer_negative_margin=0.20, + geometry_cross_view_outer_background_max_coverage=0.25, + geometry_cross_view_outer_min_views=2, + outer_uv_occupancy=False, + outer_uv_occupancy_blend_weight=0.0, + outer_uv_occupancy_gate_threshold=0.15, + outer_uv_occupancy_rescue_threshold=0.70, + outer_uv_occupancy_rescue_route_threshold=0.30, + outer_uv_component_routing=False, + outer_uv_component_seed_threshold=0.80, + outer_uv_component_grow_threshold=0.50, + outer_uv_component_min_size=2, + observed_foreground=None, + background_color_tolerance=SOLID_BACKGROUND_COLOR_TOLERANCE, + color_background_tolerance=None, + color_foreground_inset=1, + reject_semantic_fallback=False, + reject_inner_semantic_fallback=False, + include_rejected_context=True, + rejected_context_confidence_threshold=0.35, + rejected_context_margin_threshold=0.10, + rejected_context_alpha_confidence_threshold=0.50, + rejected_context_alpha_margin_threshold=0.10, + include_confidence=False, + return_details=False, + observed_color_support=None, + head_color_separation_inset=0, +): + """Route parser outputs to UV, using static mappings for affine-parser checkpoints.""" + if not 0.0 <= route_confidence_threshold <= 1.0: + raise ValueError("route_confidence_threshold must be in [0, 1].") + if not 0.0 <= route_margin_threshold <= 1.0: + raise ValueError("route_margin_threshold must be in [0, 1].") + if outer_route_confidence_threshold is None: + outer_route_confidence_threshold = route_confidence_threshold + if outer_route_margin_threshold is None: + outer_route_margin_threshold = route_margin_threshold + if not 0.0 <= outer_route_confidence_threshold <= 1.0: + raise ValueError("outer_route_confidence_threshold must be in [0, 1].") + if not 0.0 <= outer_route_margin_threshold <= 1.0: + raise ValueError("outer_route_margin_threshold must be in [0, 1].") + if not 0.0 <= outer_uv_min_coverage <= 1.0: + raise ValueError("outer_uv_min_coverage must be in [0, 1].") + if outer_uv_min_source_pixels < 1: + raise ValueError("outer_uv_min_source_pixels must be positive.") + if not 0.0 <= outer_silhouette_min_coverage <= 1.0: + raise ValueError("outer_silhouette_min_coverage must be in [0, 1].") + if outer_silhouette_dilation < 0: + raise ValueError("outer_silhouette_dilation must be non-negative.") + if outer_silhouette_min_pixels < 1: + raise ValueError("outer_silhouette_min_pixels must be positive.") + if not ( + 0.0 + <= outer_uv_component_grow_threshold + <= outer_uv_component_seed_threshold + <= 1.0 + ): + raise ValueError( + "Require outer component grow <= seed threshold in [0, 1]." + ) + if outer_uv_component_min_size < 1: + raise ValueError("outer_uv_component_min_size must be positive.") + if not 0.0 <= outer_rescue_confidence_threshold <= 1.0: + raise ValueError("outer_rescue_confidence_threshold must be in [0, 1].") + if not 0.0 <= outer_rescue_margin_threshold <= 1.0: + raise ValueError("outer_rescue_margin_threshold must be in [0, 1].") + if not 0.0 <= outer_rescue_min_coverage <= 1.0: + raise ValueError("outer_rescue_min_coverage must be in [0, 1].") + if not 0.0 <= outer_semantic_presence_threshold <= 1.0: + raise ValueError("outer_semantic_presence_threshold must be in [0, 1].") + if not 0.0 <= outer_semantic_coverage_threshold <= 1.0: + raise ValueError("outer_semantic_coverage_threshold must be in [0, 1].") + for name, value in ( + ( + "geometry_route_texel_consensus_weight", + geometry_route_texel_consensus_weight, + ), + ( + "geometry_route_preserve_outer_confidence", + geometry_route_preserve_outer_confidence, + ), + ( + "geometry_route_preserve_outer_margin", + geometry_route_preserve_outer_margin, + ), + ( + "geometry_route_consensus_outer_confidence", + geometry_route_consensus_outer_confidence, + ), + ( + "geometry_route_consensus_outer_margin", + geometry_route_consensus_outer_margin, + ), + ( + "geometry_cross_view_outer_weight", + geometry_cross_view_outer_weight, + ), + ( + "geometry_cross_view_outer_positive_confidence", + geometry_cross_view_outer_positive_confidence, + ), + ( + "geometry_cross_view_outer_positive_margin", + geometry_cross_view_outer_positive_margin, + ), + ( + "geometry_cross_view_outer_negative_confidence", + geometry_cross_view_outer_negative_confidence, + ), + ( + "geometry_cross_view_outer_negative_margin", + geometry_cross_view_outer_negative_margin, + ), + ( + "geometry_cross_view_outer_background_max_coverage", + geometry_cross_view_outer_background_max_coverage, + ), + ( + "outer_uv_occupancy_blend_weight", + outer_uv_occupancy_blend_weight, + ), + ( + "outer_uv_occupancy_gate_threshold", + outer_uv_occupancy_gate_threshold, + ), + ( + "outer_uv_occupancy_rescue_threshold", + outer_uv_occupancy_rescue_threshold, + ), + ( + "outer_uv_occupancy_rescue_route_threshold", + outer_uv_occupancy_rescue_route_threshold, + ), + ): + if not 0.0 <= float(value) <= 1.0: + raise ValueError(f"{name} must be in [0, 1].") + if not 0.0 <= rejected_context_confidence_threshold <= 1.0: + raise ValueError("rejected_context_confidence_threshold must be in [0, 1].") + if not 0.0 <= rejected_context_margin_threshold <= 1.0: + raise ValueError("rejected_context_margin_threshold must be in [0, 1].") + if not 0.0 <= rejected_context_alpha_confidence_threshold <= 1.0: + raise ValueError( + "rejected_context_alpha_confidence_threshold must be in [0, 1]." + ) + if not 0.0 <= rejected_context_alpha_margin_threshold <= 1.0: + raise ValueError( + "rejected_context_alpha_margin_threshold must be in [0, 1]." + ) + if not 0.0 <= background_color_tolerance <= 1.0: + raise ValueError("background_color_tolerance must be in [0, 1].") + if ( + color_background_tolerance is not None + and not 0.0 <= color_background_tolerance <= 1.0 + ): + raise ValueError("color_background_tolerance must be in [0, 1].") + if color_foreground_inset < 0: + raise ValueError("color_foreground_inset must be non-negative.") + if geometry_cross_view_outer_min_views < 2: + raise ValueError("geometry_cross_view_outer_min_views must be at least 2.") + if "affine" not in outputs: + conditioning = splat_predictions_to_uv_conditioning( + rendered, + outputs, + group_size=group_size, + fg_threshold=fg_threshold, + bg_color=bg_color, + ) + if return_details: + return conditioning, {"rendered": rendered, "outputs": outputs, "routing": None, "alignment": None} + return conditioning + + if renderer is None or views is None: + raise ValueError("Affine parser routing requires both renderer and views.") + views = parse_views(views) + if group_size != len(views): + raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") + + if observed_foreground is None: + observed_foreground = estimate_solid_background_foreground( + rendered, + color_tolerance=background_color_tolerance, + ) + else: + if observed_foreground.shape != rendered.shape[:1] + rendered.shape[-2:]: + raise ValueError( + "observed_foreground must have shape " + f"{rendered.shape[:1] + rendered.shape[-2:]}, got {tuple(observed_foreground.shape)}." + ) + observed_foreground = observed_foreground.to( + device=rendered.device, dtype=torch.bool + ) + source_background_rgb, _ = estimate_solid_background_color(rendered) + source_fill = rendered.new_ones(rendered.shape[0], rendered.shape[1], 1, 1) + source_fill[:, :3] = source_background_rgb.to(dtype=rendered.dtype) + + alignment = None + routing_outputs = outputs + if affine_refine: + refined_affine, alignment = refine_parser_affine( + outputs, + renderer, + views, + translation_radius_px=affine_refine_translation_px, + scale_radius=affine_refine_scale, + observed_foreground=observed_foreground, + ) + routing_outputs = dict(outputs) + routing_outputs["affine"] = refined_affine + + canonical_rendered = canonicalize_parser_render( + rendered, + routing_outputs, + mode="nearest", + fill_color=source_fill, + ) + canonical_observed_foreground = canonicalize_foreground_mask( + observed_foreground, + routing_outputs["affine"], + ) + canonical_observed_foreground_nearest = canonicalize_tensor( + observed_foreground.unsqueeze(1).to(dtype=rendered.dtype), + routing_outputs["affine"], + mode="nearest", + )[:, 0] > 0.5 + canonical_foreground_coverage_rescued = ( + canonical_observed_foreground + & ~canonical_observed_foreground_nearest + ) + color_support = build_color_sampling_support( + canonical_rendered, + canonical_observed_foreground, + source_background_rgb, + background_tolerance=color_background_tolerance, + foreground_inset=color_foreground_inset, + ) + if observed_color_support is not None: + if observed_color_support.shape != observed_foreground.shape: + raise ValueError('Observed color support must match observed foreground') + safe_source=canonicalize_tensor(observed_color_support[:,None].to(rendered.device,rendered.dtype),routing_outputs['affine'],mode='nearest')[:,0]>.5 + color_support['rejected'] |= color_support['valid'] & ~safe_source + color_support['valid'] &= safe_source + canonical_outputs = canonicalize_parser_outputs(routing_outputs) + if "surface" in canonical_outputs and "part" not in canonical_outputs: + routing = _routing_from_geometry_surface_outputs( + renderer, + views, + canonical_outputs, + fg_threshold=fg_threshold, + texel_consensus=geometry_route_texel_consensus, + texel_consensus_weight=geometry_route_texel_consensus_weight, + preserve_outer_confidence=geometry_route_preserve_outer_confidence, + preserve_outer_margin=geometry_route_preserve_outer_margin, + consensus_outer_confidence=geometry_route_consensus_outer_confidence, + consensus_outer_margin=geometry_route_consensus_outer_margin, + outer_uv_occupancy=outer_uv_occupancy, + outer_uv_occupancy_blend_weight=( + outer_uv_occupancy_blend_weight + ), + outer_uv_occupancy_gate_threshold=( + outer_uv_occupancy_gate_threshold + ), + outer_uv_occupancy_rescue_threshold=( + outer_uv_occupancy_rescue_threshold + ), + outer_uv_occupancy_rescue_route_threshold=( + outer_uv_occupancy_rescue_route_threshold + ), + outer_uv_component_routing=outer_uv_component_routing, + outer_uv_component_seed_threshold=( + outer_uv_component_seed_threshold + ), + outer_uv_component_grow_threshold=( + outer_uv_component_grow_threshold + ), + outer_uv_component_min_size=outer_uv_component_min_size, + observed_foreground=canonical_observed_foreground, + cross_view_outer_consistency=( + geometry_cross_view_outer_consistency + ), + cross_view_outer_weight=geometry_cross_view_outer_weight, + cross_view_outer_positive_confidence=( + geometry_cross_view_outer_positive_confidence + ), + cross_view_outer_positive_margin=( + geometry_cross_view_outer_positive_margin + ), + cross_view_outer_negative_confidence=( + geometry_cross_view_outer_negative_confidence + ), + cross_view_outer_negative_margin=( + geometry_cross_view_outer_negative_margin + ), + cross_view_outer_background_max_coverage=( + geometry_cross_view_outer_background_max_coverage + ), + cross_view_outer_min_views=geometry_cross_view_outer_min_views, + ) + elif "surface" not in canonical_outputs and "part" not in canonical_outputs: + routing = _routing_from_geometry_outputs( + renderer, + views, + canonical_outputs, + fg_threshold=fg_threshold, + outer_threshold=geometry_outer_threshold, + texel_consensus=geometry_route_texel_consensus, + texel_consensus_weight=geometry_route_texel_consensus_weight, + preserve_outer_confidence=geometry_route_preserve_outer_confidence, + preserve_outer_margin=geometry_route_preserve_outer_margin, + consensus_outer_confidence=geometry_route_consensus_outer_confidence, + consensus_outer_margin=geometry_route_consensus_outer_margin, + outer_uv_occupancy=outer_uv_occupancy, + outer_uv_occupancy_blend_weight=( + outer_uv_occupancy_blend_weight + ), + outer_uv_occupancy_gate_threshold=( + outer_uv_occupancy_gate_threshold + ), + outer_uv_occupancy_rescue_threshold=( + outer_uv_occupancy_rescue_threshold + ), + outer_uv_occupancy_rescue_route_threshold=( + outer_uv_occupancy_rescue_route_threshold + ), + outer_uv_component_routing=outer_uv_component_routing, + outer_uv_component_seed_threshold=( + outer_uv_component_seed_threshold + ), + outer_uv_component_grow_threshold=( + outer_uv_component_grow_threshold + ), + outer_uv_component_min_size=outer_uv_component_min_size, + observed_foreground=canonical_observed_foreground, + cross_view_outer_consistency=( + geometry_cross_view_outer_consistency + ), + cross_view_outer_weight=geometry_cross_view_outer_weight, + cross_view_outer_positive_confidence=( + geometry_cross_view_outer_positive_confidence + ), + cross_view_outer_positive_margin=( + geometry_cross_view_outer_positive_margin + ), + cross_view_outer_negative_confidence=( + geometry_cross_view_outer_negative_confidence + ), + cross_view_outer_negative_margin=( + geometry_cross_view_outer_negative_margin + ), + cross_view_outer_background_max_coverage=( + geometry_cross_view_outer_background_max_coverage + ), + cross_view_outer_min_views=geometry_cross_view_outer_min_views, + ) + else: + routing = _routing_from_affine_outputs( + renderer, + views, + canonical_outputs, + fg_threshold=fg_threshold, + semantic_gate=semantic_gate, + ) + from SkingToolkit.dense_uv_parser.accessories import apply_accessory_routing + accessory_supported = apply_accessory_routing( + routing, canonical_outputs, canonical_observed_foreground, renderer, views) + raw_foreground = routing["foreground"] + routing["secondary"] = routing.get( + "secondary", torch.zeros_like(raw_foreground) + ) & canonical_observed_foreground + selected_outer = routing["layer"] == 1 + confidence_threshold = torch.where( + selected_outer, + routing["confidence"].new_tensor(outer_route_confidence_threshold), + routing["confidence"].new_tensor(route_confidence_threshold), + ) + margin_threshold = torch.where( + selected_outer, + routing["confidence_margin_ratio"].new_tensor(outer_route_margin_threshold), + routing["confidence_margin_ratio"].new_tensor(route_margin_threshold), + ) + strict_trusted = ( + raw_foreground + & canonical_observed_foreground + & (routing["confidence"] >= confidence_threshold) + & (routing["confidence_margin_ratio"] >= margin_threshold) + ) + outer_consensus_gate_rejected = routing.get( + "consensus_outer_gate_rejected", + torch.zeros_like(selected_outer), + ) + # A deferred outer route is intentionally unknown. It may be restored by + # exact geometry or atlas occupancy evidence below, but broad semantic part + # presence must not turn the rejected local observation back into outer. + deferred_outer_gate_rejected = ( + selected_outer & outer_consensus_gate_rejected + ) + strict_trusted = strict_trusted & ~deferred_outer_gate_rejected + if reject_semantic_fallback: + rejected_fallback = routing["semantic_fallback"] & ( + selected_outer | reject_inner_semantic_fallback + ) + strict_trusted = strict_trusted & ~rejected_fallback + geometry_supported_outer = torch.zeros_like(selected_outer) + semantic_supported_outer = torch.zeros_like(selected_outer) + occupancy_supported_outer = torch.zeros_like(selected_outer) + outer_geometry_rescued = torch.zeros_like(selected_outer) + outer_semantic_rescued = torch.zeros_like(selected_outer) + outer_occupancy_rescued = torch.zeros_like(selected_outer) + if outer_geometry_rescue: + geometry_supported_outer = _geometry_supported_outer_texels( + routing, + renderer, + views, + ) + rescue_trusted = ( + raw_foreground + & canonical_observed_foreground + & geometry_supported_outer + & (routing["confidence"] >= outer_rescue_confidence_threshold) + & ( + routing["confidence_margin_ratio"] + >= outer_rescue_margin_threshold + ) + ) + if reject_semantic_fallback: + rescue_trusted = rescue_trusted & ~rejected_fallback + outer_geometry_rescued = rescue_trusted & ~strict_trusted + trusted = strict_trusted | rescue_trusted + else: + trusted = strict_trusted + if outer_semantic_rescue: + semantic_supported_outer = _semantic_supported_outer_pixels( + canonical_outputs, + routing, + group_size, + presence_threshold=outer_semantic_presence_threshold, + coverage_threshold=outer_semantic_coverage_threshold, + ) + semantic_rescue_trusted = ( + raw_foreground + & canonical_observed_foreground + & semantic_supported_outer + & ~deferred_outer_gate_rejected + & (routing["confidence"] >= outer_rescue_confidence_threshold) + & ( + routing["confidence_margin_ratio"] + >= outer_rescue_margin_threshold + ) + ) + if reject_semantic_fallback: + semantic_rescue_trusted = semantic_rescue_trusted & ~rejected_fallback + outer_semantic_rescued = semantic_rescue_trusted & ~trusted + trusted = trusted | semantic_rescue_trusted + occupancy_available = routing.get( + "outer_uv_occupancy_available", + torch.zeros_like(selected_outer), + ) + occupancy_supported_outer = ( + selected_outer + & occupancy_available + & ( + routing.get( + "projected_outer_occupancy", + torch.zeros_like(routing["confidence"]), + ) + >= float(outer_uv_occupancy_rescue_threshold) + ) + ) + if outer_uv_occupancy and occupancy_supported_outer.any(): + occupancy_rescue_trusted = ( + raw_foreground + & canonical_observed_foreground + & occupancy_supported_outer + & (routing["confidence"] >= outer_rescue_confidence_threshold) + & ( + routing["confidence_margin_ratio"] + >= outer_rescue_margin_threshold + ) + ) + if reject_semantic_fallback: + occupancy_rescue_trusted = ( + occupancy_rescue_trusted & ~rejected_fallback + ) + outer_occupancy_rescued = occupancy_rescue_trusted & ~trusted + trusted = trusted | occupancy_rescue_trusted + trusted = trusted | accessory_supported + outer_silhouette_coverage = torch.ones_like(routing["confidence"]) + outer_silhouette_assessed = torch.zeros_like(selected_outer) + outer_silhouette_rejected = torch.zeros_like(selected_outer) + if outer_silhouette_consistency: + ( + outer_silhouette_coverage, + outer_silhouette_assessed, + ) = _outer_silhouette_coverage( + # A one-pixel antialias/background fringe often survives flood + # fill. It is unsafe silhouette evidence even when it is already + # excluded from RGB pickup. Use the inset foreground here so that + # such a fringe cannot validate a false outer shell. + color_support["interior"], + routing["flat_uv"], + renderer, + views, + dilation=outer_silhouette_dilation, + min_pixels=outer_silhouette_min_pixels, + ) + outer_silhouette_rejected = ( + trusted + & selected_outer + & outer_silhouette_assessed + & ( + outer_silhouette_coverage + < float(outer_silhouette_min_coverage) + ) + ) + # Thin crown teeth need not fill most of a projected square. A learned + # whole-object decision plus confident pixel semantics may preserve them; + # foreground and valid outer geometry are still mandatory. All other + # objects retain the established silhouette veto. + crown_rescue = torch.zeros_like(trusted) + if 'headwear_presence_logits' in canonical_outputs and 'headwear_supported' in routing: + crown_rescue = (routing['headwear_supported'] & (routing['headwear_family']==4) + & (routing['headwear_probability']>=.8) + & canonical_observed_foreground & selected_outer) + routing['headwear_silhouette_rescued'] = outer_silhouette_rejected & crown_rescue + outer_silhouette_rejected &= ~crown_rescue + trusted = trusted & ~outer_silhouette_rejected + outer_uv_coverage = torch.ones_like(routing["confidence"]) + outer_required_coverage = torch.zeros_like(routing["confidence"]) + if outer_uv_min_coverage > 0.0: + views_per_group = len(views) + group_count = trusted.shape[0] // views_per_group + if "surface" in canonical_outputs: + pixel_coverage = _surface_aware_outer_coverage( + routing, + trusted, + selected_outer, + renderer, + views, + ) + else: + expected = routing["confidence"].new_zeros(UV_SIZE * UV_SIZE) + for view in views: + static = build_static_surface_routing(renderer, view, trusted.device) + expected_uv = static["flat_uv"][1][static["masks"][1]] + expected.scatter_add_( + 0, + expected_uv, + torch.ones(expected_uv.shape[0], device=trusted.device, dtype=expected.dtype), + ) + + item_groups = torch.arange(trusted.shape[0], device=trusted.device) // views_per_group + grouped_uv = routing["flat_uv"] + item_groups.view(-1, 1, 1) * (UV_SIZE * UV_SIZE) + observed_mask = trusted & selected_outer + observed = routing["confidence"].new_zeros(group_count * UV_SIZE * UV_SIZE) + observed_uv = grouped_uv[observed_mask] + observed.scatter_add_( + 0, + observed_uv, + torch.ones(observed_uv.shape[0], device=trusted.device, dtype=observed.dtype), + ) + coverage = observed.reshape(group_count, UV_SIZE * UV_SIZE) / expected.clamp_min(1.0) + pixel_coverage = coverage[item_groups.view(-1, 1, 1), routing["flat_uv"]] + outer_uv_coverage = torch.where(selected_outer, pixel_coverage, outer_uv_coverage) + required_coverage = torch.full_like(pixel_coverage, outer_uv_min_coverage) + if ( + outer_geometry_rescue + or outer_semantic_rescue + or outer_uv_occupancy + ): + required_coverage = torch.where( + geometry_supported_outer + | semantic_supported_outer + | occupancy_supported_outer, + required_coverage.new_tensor( + min(outer_uv_min_coverage, outer_rescue_min_coverage) + ), + required_coverage, + ) + outer_required_coverage = torch.where( + selected_outer, + required_coverage, + outer_required_coverage, + ) + trusted = trusted & (~selected_outer | (pixel_coverage >= required_coverage) | accessory_supported) + outer_source_support = torch.zeros_like(routing["flat_uv"]) + outer_source_rejected = torch.zeros_like(selected_outer) + if outer_uv_min_source_pixels > 1: + views_per_group = len(views) + group_count = trusted.shape[0] // views_per_group + item_groups = ( + torch.arange(trusted.shape[0], device=trusted.device) + // views_per_group + ) + grouped_uv = routing["flat_uv"] + item_groups.view(-1, 1, 1) * ( + UV_SIZE * UV_SIZE + ) + outer_candidates = trusted & selected_outer + source_counts = torch.zeros( + group_count * UV_SIZE * UV_SIZE, + dtype=torch.long, + device=trusted.device, + ) + source_counts.scatter_add_( + 0, + grouped_uv[outer_candidates], + torch.ones_like(grouped_uv[outer_candidates]), + ) + outer_source_support = source_counts[grouped_uv] + outer_source_rejected = ( + outer_candidates + & (outer_source_support < int(outer_uv_min_source_pixels)) + & ~accessory_supported + ) + trusted = trusted & ~outer_source_rejected + routing["raw_foreground"] = raw_foreground + routing["observed_foreground"] = canonical_observed_foreground + routing["canonical_foreground_coverage_rescued"] = ( + canonical_foreground_coverage_rescued + ) + routing["background_rejected"] = raw_foreground & ~canonical_observed_foreground + routing["rejected"] = raw_foreground & ~trusted + routing["outer_geometry_supported"] = geometry_supported_outer + routing["outer_geometry_rescued"] = outer_geometry_rescued & trusted + routing["outer_semantic_supported"] = semantic_supported_outer + routing["outer_semantic_rescued"] = outer_semantic_rescued & trusted + routing["outer_occupancy_supported"] = occupancy_supported_outer + routing["outer_occupancy_rescued"] = outer_occupancy_rescued & trusted + routing["outer_required_coverage"] = outer_required_coverage + routing["outer_silhouette_coverage"] = outer_silhouette_coverage + routing["outer_silhouette_assessed"] = outer_silhouette_assessed + routing["outer_silhouette_rejected"] = outer_silhouette_rejected + routing["outer_source_support"] = outer_source_support + routing["outer_source_rejected"] = outer_source_rejected + routing["foreground"] = trusted + color_foreground = trusted & color_support["valid"] + boundary_excluded = torch.zeros_like(trusted) + if head_color_separation_inset > 0: + # A layer boundary is also an RGB mixing boundary. Keep occupancy, but + # do not let a thin contaminated inner fringe become an inpainting seed. + radius=int(head_color_separation_inset) + head_outer=trusted&(routing['layer']==1)&(routing['part']==0) + near_outer=F.max_pool2d(head_outer[:,None].float(),2*radius+1,1,radius)[:,0]>.5 + boundary_excluded=near_outer&(routing['layer']==0)&(routing['part']==0) + phone=routing.get('headphone_colour_conflict') + if phone is not None: + # Compare a semantic accessory observation with nearby accepted outer + # geometry; exclude contradictory RGB without inventing outer texels. + nearby=F.max_pool2d(head_outer[:,None].float(),17,1,8)[:,0]>.5 + conflict=phone&nearby&(routing['layer']==0)&(routing['part']==0) + boundary_excluded|=conflict + routing['headphone_colour_rejected']=conflict + color_foreground &= ~boundary_excluded + headwear_layer = routing.get('headwear_layer') + if headwear_layer is not None: + # A component outside the chosen layer's projected support is unknown; + # sampling it into the other layer would duplicate its material. + conflict = (headwear_layer >= 0) & (routing['layer'] != headwear_layer) & (routing['part'] == 0) + color_foreground &= ~conflict + family = routing['headwear_family'] + boundary = torch.zeros_like(trusted) + for component in (1,2,3,4): + region = (family == component) & routing['headwear_supported'] + interior = -F.max_pool2d(-region[:,None].float(),3,1,1)[:,0] > .5 + boundary |= region & ~interior + color_foreground &= ~boundary + routing['headwear_color_rejected'] = conflict | boundary + routing['head_color_boundary_excluded']=boundary_excluded + routing["color_foreground"] = color_foreground + routing["color_interior"] = color_support["interior"] + routing["color_background_like"] = color_support["background_like"] + routing["color_rejected"] = trusted & color_support["rejected"] + if "secondary_routed" in routing: + routing["secondary_rejected"] = routing["secondary"] & ~trusted + routing["secondary_routed"] = routing["secondary"] & trusted + routing["outer_uv_coverage"] = outer_uv_coverage + conditioning = splat_to_uv_conditioning( + canonical_rendered, + color_foreground, + routing["layer"], + routing["flat_uv"], + group_size=group_size, + bg_color=bg_color, + confidence=routing["confidence"], + sampling_quality=( + routing.get("texel_center_score") + + 0.05 * color_support["interior"].to( + dtype=routing["confidence"].dtype + ) + ), + color_aggregation=color_aggregation, + include_confidence=include_confidence, + ) + rejected_context = ( + raw_foreground + & canonical_observed_foreground + & ~trusted + & color_support["valid"] + & (routing["confidence"] >= rejected_context_confidence_threshold) + & ( + routing["confidence_margin_ratio"] + >= rejected_context_margin_threshold + ) + ) + if reject_semantic_fallback: + rejected_context = rejected_context & ~rejected_fallback + rejected_context_alpha_supported = ( + rejected_context + & selected_outer + & (geometry_supported_outer | semantic_supported_outer) + & ( + routing["confidence"] + >= rejected_context_alpha_confidence_threshold + ) + & ( + routing["confidence_margin_ratio"] + >= rejected_context_alpha_margin_threshold + ) + ) + context_alpha_rescue_uv = conditioning.new_zeros( + conditioning.shape[0], 1, UV_SIZE, UV_SIZE + ) + if include_rejected_context and include_confidence and rejected_context.any(): + context = splat_to_uv_conditioning( + canonical_rendered, + rejected_context, + routing["layer"], + routing["flat_uv"], + group_size=group_size, + bg_color=bg_color, + confidence=routing["confidence"], + sampling_quality=( + routing.get("texel_center_score") + + 0.05 * color_support["interior"].to( + dtype=routing["confidence"].dtype + ) + ), + color_aggregation=color_aggregation, + include_confidence=True, + ) + for offset in (0, 6): + trusted_evidence = conditioning[:, offset + 4 : offset + 5] > 0.5 + context_evidence = context[:, offset + 4 : offset + 5] > 0.5 + fill = ~trusted_evidence & context_evidence + conditioning[:, offset : offset + 4] = torch.where( + fill.expand(-1, 4, -1, -1), + context[:, offset : offset + 4], + conditioning[:, offset : offset + 4], + ) + conditioning[:, offset + 5 : offset + 6] = torch.where( + fill, + context[:, offset + 5 : offset + 6], + conditioning[:, offset + 5 : offset + 6], + ) + if rejected_context_alpha_supported.any(): + alpha_context = splat_to_uv_conditioning( + canonical_rendered, + rejected_context_alpha_supported, + routing["layer"], + routing["flat_uv"], + group_size=group_size, + bg_color=bg_color, + confidence=routing["confidence"], + sampling_quality=( + routing.get("texel_center_score") + + 0.05 * color_support["interior"].to( + dtype=routing["confidence"].dtype + ) + ), + color_aggregation=color_aggregation, + include_confidence=True, + ) + context_alpha_rescue_uv = alpha_context[:, 10:11] > 0.5 + routing["rejected_context"] = rejected_context + routing["rejected_context_alpha_supported"] = ( + rejected_context_alpha_supported + ) + if return_details: + return conditioning, { + "rendered": canonical_rendered, + "color_source_support": color_support["valid"], + "outputs": canonical_outputs, + "routing": routing, + "alignment": alignment, + "context_alpha_rescue_uv": context_alpha_rescue_uv, + } + return conditioning + + +def splat_deterministic_targets_to_uv_conditioning( + rendered, + targets, + renderer, + views, + group_size=1, + bg_color=(128, 128, 128), + color_aggregation="grid_mode", +): + """Splat ground-truth labels through the same fixed mapping used by affine mode.""" + views = parse_views(views) + if group_size != len(views): + raise ValueError(f"group_size={group_size} must equal the number of views ({len(views)}).") + canonical_rendered = canonicalize_tensor(rendered, targets["affine"], mode="nearest") + canonical_targets = canonicalize_dense_targets(targets) + requested_surface = canonical_targets["surface"] + N, H, W = requested_surface.shape + fg = torch.zeros_like(requested_surface, dtype=torch.bool) + routed_layer = torch.zeros_like(requested_surface) + flat_uv = torch.zeros_like(requested_surface) + texel_center_score = canonical_rendered.new_zeros(requested_surface.shape) + + for view_index, view in enumerate(views): + selection = slice(view_index, N, len(views)) + static = build_static_surface_routing(renderer, view, requested_surface.device) + if static["masks"].shape[-2:] != (H, W): + raise ValueError( + f"View {view!r} mapping shape {tuple(static['masks'].shape[-2:])} does not match target shape {(H, W)}." + ) + routed = _select_static_surface(static, requested_surface[selection]) + fg[selection] = ( + (canonical_targets["foreground"][selection, 0] > 0.5) + & (requested_surface[selection] != IGNORE_INDEX) + & routed["valid"] + ) + routed_layer[selection] = routed["layer"] + flat_uv[selection] = routed["flat_uv"] + texel_center_score[selection] = routed["texel_center_score"] + + return splat_to_uv_conditioning( + canonical_rendered, + fg, + routed_layer, + flat_uv, + group_size=group_size, + bg_color=bg_color, + sampling_quality=texel_center_score, + color_aggregation=color_aggregation, + ) + + +def splat_predictions_to_uv_conditioning( + rendered, + outputs, + group_size=1, + fg_threshold=0.5, + bg_color=(128, 128, 128), +): + """Splat parser predictions into the parser's layered UV evidence layout.""" + foreground_prob = torch.sigmoid(outputs["foreground"])[:, 0] + fg = foreground_prob > fg_threshold + fg = fg & (rendered[:, 3] > 1e-4) + layer_prob = torch.softmax(outputs["layer"], dim=1) + layer_confidence, layer = layer_prob.max(dim=1) + flat_uv = prediction_flat_uv(outputs) + confidence = foreground_prob * layer_confidence + if "uv_x" in outputs and "uv_y" in outputs: + uv_x_confidence = torch.softmax(outputs["uv_x"], dim=1).amax(dim=1) + uv_y_confidence = torch.softmax(outputs["uv_y"], dim=1).amax(dim=1) + confidence = confidence * torch.sqrt((uv_x_confidence * uv_y_confidence).clamp_min(1e-8)) + + return splat_to_uv_conditioning( + rendered, + fg, + layer, + flat_uv, + group_size=group_size, + bg_color=bg_color, + confidence=confidence, + ) + + +def prediction_flat_uv(outputs): + if "uv_x" in outputs and "uv_y" in outputs: + x = outputs["uv_x"].argmax(dim=1).clamp(0, UV_SIZE - 1) + y = outputs["uv_y"].argmax(dim=1).clamp(0, UV_SIZE - 1) + else: + uv = outputs["uv"].clamp(0.0, 1.0) + x = (uv[:, 0] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) + y = (uv[:, 1] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) + return y * UV_SIZE + x + + +def prediction_uv01(outputs): + if "uv_x" in outputs and "uv_y" in outputs: + x = outputs["uv_x"].argmax(dim=1).to(dtype=outputs["uv"].dtype) / (UV_SIZE - 1) + y = outputs["uv_y"].argmax(dim=1).to(dtype=outputs["uv"].dtype) / (UV_SIZE - 1) + return torch.stack([x, y], dim=1) + return outputs["uv"].clamp(0.0, 1.0) + + +def splat_targets_to_uv_conditioning(rendered, targets, group_size=1, bg_color=(128, 128, 128)): + if rendered.dim() != 4: + raise ValueError(f"Expected rendered tensor as NCHW, got {tuple(rendered.shape)}.") + layer = targets["layer"] + fg = (targets["foreground"][:, 0] > 0.5) & (layer != IGNORE_INDEX) & (rendered[:, 3] > 1e-4) + uv = targets["uv"].clamp(0.0, 1.0) + x = (uv[:, 0] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) + y = (uv[:, 1] * (UV_SIZE - 1)).round().long().clamp(0, UV_SIZE - 1) + flat_uv = y * UV_SIZE + x + safe_layer = torch.where(layer == IGNORE_INDEX, torch.zeros_like(layer), layer) + return splat_to_uv_conditioning(rendered, fg, safe_layer, flat_uv, group_size=group_size, bg_color=bg_color) + + +def _select_exact_mode_candidates( + values, + target_uv, + scores, + color_tie_break=None, +): + """Select one real source pixel from the most frequent 8-bit RGB value per UV texel.""" + if target_uv.numel() == 0: + return values, target_uv, scores + + rgb8 = (values[:3].clamp(0.0, 1.0) * 255.0).round().long() + color_code = rgb8[0] | (rgb8[1] << 8) | (rgb8[2] << 16) + color_count = 1 << 24 + combined_key = target_uv.long() * color_count + color_code + unique_keys, inverse = torch.unique(combined_key, sorted=False, return_inverse=True) + votes = scores.new_zeros(unique_keys.shape[0]) + votes.index_add_(0, inverse, torch.ones_like(scores)) + tie_break_votes = None + if color_tie_break is not None: + tie_break_votes = scores.new_zeros(unique_keys.shape[0]) + tie_break_votes.index_add_(0, inverse, color_tie_break.float()) + + positions = torch.arange(target_uv.shape[0], device=target_uv.device) + first_position = torch.full( + (unique_keys.shape[0],), + target_uv.shape[0], + dtype=torch.long, + device=target_uv.device, + ) + first_position.scatter_reduce_(0, inverse, positions, reduce="amin", include_self=True) + + unique_uv = unique_keys.div(color_count, rounding_mode="floor") + best_votes = votes.new_full((UV_SIZE * UV_SIZE,), -torch.inf) + best_votes.scatter_reduce_(0, unique_uv, votes, reduce="amax", include_self=True) + tied_winner = votes >= (best_votes[unique_uv] - 1e-7) + if tie_break_votes is not None: + best_tie_break = tie_break_votes.new_full( + (UV_SIZE * UV_SIZE,), -torch.inf + ) + best_tie_break.scatter_reduce_( + 0, + unique_uv[tied_winner], + tie_break_votes[tied_winner], + reduce="amax", + include_self=True, + ) + tied_winner = tied_winner & ( + tie_break_votes >= (best_tie_break[unique_uv] - 1e-7) + ) + earliest_winner = torch.full( + (UV_SIZE * UV_SIZE,), + target_uv.shape[0], + dtype=torch.long, + device=target_uv.device, + ) + earliest_winner.scatter_reduce_( + 0, + unique_uv[tied_winner], + first_position[tied_winner], + reduce="amin", + include_self=True, + ) + winning_unique = tied_winner & (first_position == earliest_winner[unique_uv]) + winning_key = torch.full( + (UV_SIZE * UV_SIZE,), + -1, + dtype=torch.long, + device=target_uv.device, + ) + winning_key[unique_uv[winning_unique]] = unique_keys[winning_unique] + in_winning_color = combined_key == winning_key[target_uv] + + values = values[:, in_winning_color] + target_uv = target_uv[in_winning_color] + scores = scores[in_winning_color] + best_scores = scores.new_full((UV_SIZE * UV_SIZE,), -torch.inf) + best_scores.scatter_reduce_(0, target_uv, scores, reduce="amax", include_self=True) + tied_representative = scores >= (best_scores[target_uv] - 1e-7) + positions = torch.arange(target_uv.shape[0], device=target_uv.device) + first_representative = torch.full( + (UV_SIZE * UV_SIZE,), + target_uv.shape[0], + dtype=torch.long, + device=target_uv.device, + ) + first_representative.scatter_reduce_( + 0, + target_uv[tied_representative], + positions[tied_representative], + reduce="amin", + include_self=True, + ) + selected = tied_representative & (positions == first_representative[target_uv]) + return values[:, selected], target_uv[selected], scores[selected] + + +def splat_to_uv_conditioning( + rendered, + fg, + layer, + flat_uv, + group_size=1, + bg_color=(128, 128, 128), + confidence=None, + sampling_quality=None, + color_aggregation="grid_mode", + include_confidence=False, +): + if rendered.dim() != 4: + raise ValueError(f"Expected rendered tensor as NCHW, got {tuple(rendered.shape)}.") + N, _, _, _ = rendered.shape + if N % group_size != 0: + raise ValueError(f"N={N} must be divisible by group_size={group_size}.") + if color_aggregation not in SPLAT_COLOR_AGGREGATIONS: + raise ValueError( + f"Unknown color_aggregation={color_aggregation!r}; " + f"expected one of {SPLAT_COLOR_AGGREGATIONS}." + ) + + groups = N // group_size + device = rendered.device + dtype = rendered.dtype + accum = rendered.new_zeros(groups, LAYER_CLASSES, 4, UV_SIZE * UV_SIZE) + counts = rendered.new_zeros(groups, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE) + confidence_map = rendered.new_zeros( + groups, LAYER_CLASSES, 1, UV_SIZE * UV_SIZE + ) + + if confidence is None: + confidence = rendered.new_ones(fg.shape) + select_highest_confidence = False + else: + select_highest_confidence = True + if sampling_quality is None: + sampling_quality = rendered.new_ones(fg.shape) + elif sampling_quality.shape != fg.shape: + raise ValueError("sampling_quality must have the same shape as fg.") + + for group in range(groups): + group_start = group * group_size + group_end = group_start + group_size + for layer_index in range(LAYER_CLASSES): + candidate_values = [] + candidate_uv = [] + candidate_confidence = [] + candidate_sampling_quality = [] + for item in range(group_start, group_end): + item_mask = fg[item] & (layer[item] == layer_index) + if not item_mask.any(): + continue + candidate_values.append(rendered[item, :, item_mask]) + candidate_uv.append(flat_uv[item, item_mask]) + candidate_confidence.append(confidence[item, item_mask]) + candidate_sampling_quality.append( + sampling_quality[item, item_mask] + ) + if not candidate_uv: + continue + + values = torch.cat(candidate_values, dim=1) + target_uv = torch.cat(candidate_uv, dim=0) + scores = torch.cat(candidate_confidence, dim=0).float() + quality = torch.cat(candidate_sampling_quality, dim=0).float() + if color_aggregation == "texel_center": + # Lexicographic selection: prefer the source sample closest to + # the integer UV center, then use parser confidence only to + # break effectively equal geometric samples. This recovers an + # original texel from bilinear renders instead of choosing a + # top/left boundary blend by scan order. + selection_score = quality + scores.clamp(0.0, 1.0) * 1e-6 + best_quality = selection_score.new_full( + (UV_SIZE * UV_SIZE,), -torch.inf + ) + best_quality.scatter_reduce_( + 0, + target_uv, + selection_score, + reduce="amax", + include_self=True, + ) + is_best = selection_score >= ( + best_quality[target_uv] - 1e-7 + ) + candidate_indices = torch.arange( + target_uv.shape[0], device=device + ) + first_best = torch.full( + (UV_SIZE * UV_SIZE,), + target_uv.shape[0], + dtype=torch.long, + device=device, + ) + first_best.scatter_reduce_( + 0, + target_uv[is_best], + candidate_indices[is_best], + reduce="amin", + include_self=True, + ) + selected_indices = first_best[ + first_best < target_uv.shape[0] + ] + values = values[:, selected_indices] + target_uv = target_uv[selected_indices] + scores = scores[selected_indices] + if color_aggregation == "grid_mode": + # The fitted renderer grid owns UV membership. Select the + # dominant real 8-bit RGB within each layer/UV cell; UV-center + # quality is only a tie-break between equally frequent colors. + values, target_uv, scores = _select_exact_mode_candidates( + values, + target_uv, + scores, + color_tie_break=quality, + ) + if color_aggregation == "exact_mode": + values, target_uv, scores = _select_exact_mode_candidates( + values, target_uv, scores + ) + if color_aggregation == "best" and select_highest_confidence: + best_scores = scores.new_full((UV_SIZE * UV_SIZE,), -torch.inf) + best_scores.scatter_reduce_(0, target_uv, scores, reduce="amax", include_self=True) + is_best = scores >= (best_scores[target_uv] - 1e-7) + candidate_indices = torch.arange(target_uv.shape[0], device=device) + first_best = torch.full( + (UV_SIZE * UV_SIZE,), + target_uv.shape[0], + dtype=torch.long, + device=device, + ) + first_best.scatter_reduce_( + 0, + target_uv[is_best], + candidate_indices[is_best], + reduce="amin", + include_self=True, + ) + selected_indices = first_best[first_best < target_uv.shape[0]] + values = values[:, selected_indices] + target_uv = target_uv[selected_indices] + scores = scores[selected_indices] + + accum[group, layer_index].index_add_(1, target_uv, values) + counts[group, layer_index, 0].index_add_( + 0, + target_uv, + torch.ones(target_uv.shape[0], dtype=dtype, device=device), + ) + confidence_map[group, layer_index, 0].scatter_reduce_( + 0, + target_uv, + scores.to(dtype=dtype), + reduce="amax", + include_self=True, + ) + + known = (counts > 0).to(dtype=dtype) + aggregated = accum / counts.clamp_min(1.0) + bg = rendered.new_tensor(bg_color, dtype=dtype).view(1, 1, 3, 1) / 255.0 + rgb = torch.where( + known.expand(-1, -1, 3, -1) > 0, + aggregated[:, :, :3], + bg.expand(groups, LAYER_CLASSES, 3, UV_SIZE * UV_SIZE), + ) + alpha = torch.where(known > 0, aggregated[:, :, 3:4], torch.zeros_like(aggregated[:, :, 3:4])) + layer_channels = [rgb, alpha, known] + if include_confidence: + layer_channels.append(confidence_map * known) + layers = torch.cat(layer_channels, dim=2).reshape( + groups, -1, UV_SIZE, UV_SIZE + ) + return layers.clamp(0.0, 1.0) + + +def conditioning_to_pred_uv(conditioning): + """Merge two-layer parser conditioning into a preliminary RGBA skin atlas. + + Conditioning stores either ``RGBA + evidence`` (legacy 10-channel input) + or ``RGBA + evidence + confidence`` (12-channel input) for each layer. + Their Minecraft UV rectangles normally do not overlap, but outer data wins + if both layers mark the same texel as observed. + Unknown RGB retains the conditioning background as a useful placeholder; + unknown alpha remains transparent until Minecraft base-alpha finalization. + """ + squeeze_batch = conditioning.dim() == 3 + if squeeze_batch: + conditioning = conditioning.unsqueeze(0) + if conditioning.dim() != 4 or conditioning.shape[1] not in (10, 12): + raise ValueError( + "Expected 10- or 12-channel parser conditioning, " + f"got {tuple(conditioning.shape)}." + ) + + inner_rgba = conditioning[:, 0:4] + inner_known = conditioning[:, 4:5] > 0.5 + outer_offset = 6 if conditioning.shape[1] == 12 else 5 + outer_rgba = conditioning[:, outer_offset : outer_offset + 4] + outer_known = conditioning[:, outer_offset + 4 : outer_offset + 5] > 0.5 + known = inner_known | outer_known + + rgba = torch.where(outer_known.expand_as(outer_rgba), outer_rgba, inner_rgba) + rgba[:, 3:4] = torch.where( + known, + rgba[:, 3:4], + torch.zeros_like(rgba[:, 3:4]), + ) + rgba = rgba.clamp(0.0, 1.0) + return rgba[0] if squeeze_batch else rgba + + +PART_PALETTE = ( + (239, 83, 80), + (255, 202, 40), + (102, 187, 106), + (38, 166, 154), + (66, 165, 245), + (171, 71, 188), +) +LAYER_PALETTE = ( + (72, 169, 166), + (255, 179, 71), +) +ROUTE_ROLE_PALETTE = ( + (72, 169, 166), + (255, 179, 71), + (218, 112, 232), +) +FACE_NAMES = ("front", "back", "left", "right", "top", "bottom") +FACE_PALETTE = ( + (239, 83, 80), + (171, 71, 188), + (66, 165, 245), + (38, 166, 154), + (255, 202, 40), + (117, 117, 117), +) +LAYER_FACE_PALETTE = ( + (183, 28, 28), + (106, 27, 154), + (21, 101, 192), + (0, 121, 107), + (245, 166, 35), + (84, 84, 84), + (255, 128, 125), + (218, 112, 232), + (126, 208, 255), + (102, 221, 203), + (255, 224, 130), + (210, 210, 210), +) + + +def combine_layer_face(layer, face): + """Encode inner/outer x six cube faces as 12 visualization classes.""" + valid = (layer != IGNORE_INDEX) & (face != IGNORE_INDEX) + combined = torch.full_like(layer, IGNORE_INDEX) + combined[valid] = layer[valid] * FACE_CLASSES + face[valid] + return combined + + +def bg_tensor(bg_color, reference): + return reference.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 + + +def colorize_labels(labels, palette, bg_color, reference): + N, H, W = labels.shape + bg = bg_tensor(bg_color, reference).expand(N, 3, H, W).clone() + valid = labels != IGNORE_INDEX + if not valid.any(): + return bg + palette_tensor = reference.new_tensor(palette, dtype=reference.dtype) / 255.0 + safe_labels = labels.clamp(0, len(palette) - 1) + out = bg.permute(0, 2, 3, 1) + out[valid] = palette_tensor[safe_labels[valid]] + return bg + + +def build_geometry_grid_debug(renderer, views, item_count, reference, bg_color=(128, 128, 128)): + """Build fitted inner/outer cuboid face maps with projected UV texel boundaries.""" + views = parse_views(views) + if not views: + raise ValueError("At least one view is required for geometry-grid debug output.") + colors = [[], []] + masks = [[], []] + edges = [[], []] + for item_index in range(item_count): + static = build_static_surface_routing(renderer, views[item_index % len(views)], reference.device) + for layer_index in range(LAYER_CLASSES): + mask = static["masks"][layer_index] + face = static["face"][layer_index] + flat_uv = static["flat_uv"][layer_index] + labels = combine_layer_face( + torch.full_like(face, layer_index), + torch.where(mask, face, torch.full_like(face, IGNORE_INDEX)), + ).unsqueeze(0) + color = colorize_labels(labels, LAYER_FACE_PALETTE, bg_color, reference[:1])[0] + + edge = torch.zeros_like(mask) + different_x = (flat_uv[:, 1:] != flat_uv[:, :-1]) | (face[:, 1:] != face[:, :-1]) + different_y = (flat_uv[1:, :] != flat_uv[:-1, :]) | (face[1:, :] != face[:-1, :]) + edge[:, 1:] |= mask[:, 1:] & mask[:, :-1] & different_x + edge[1:, :] |= mask[1:, :] & mask[:-1, :] & different_y + interior = -F.max_pool2d(-mask.float().unsqueeze(0).unsqueeze(0), 3, 1, 1)[0, 0] + edge |= mask & (interior < 0.5) + edge_color = reference.new_tensor((0.08, 0.08, 0.08)).view(3, 1, 1) + color = torch.where(edge.unsqueeze(0), edge_color, color) + + colors[layer_index].append(color) + masks[layer_index].append(mask) + edges[layer_index].append(edge) + + return ( + torch.stack(colors[0]), + torch.stack(colors[1]), + torch.stack(masks[0]), + torch.stack(masks[1]), + torch.stack(edges[0]), + torch.stack(edges[1]), + ) + + +def fill_geometry_grid_debug(rendered, foreground, layer, geometry_debug, bg_color=(128, 128, 128)): + """Show only source RGB pixels actually routed to each predicted layer.""" + del geometry_debug + rgb = rendered[:, :3] + bg = rgb.new_tensor(bg_color).view(1, 3, 1, 1) / 255.0 + + def fill(layer_index): + classified = (foreground & (layer == layer_index)).unsqueeze(1) + return torch.where(classified, rgb, bg.expand_as(rgb)) + + return fill(0), fill(1) + + +def overlay_geometry_grid_debug( + rendered, + geometry_debug, + tint_alpha=0.12, + base_images=None, +): + """Overlay fitted inner/outer UV texel grids on the canonical source image.""" + if not 0.0 <= tint_alpha <= 1.0: + raise ValueError(f"tint_alpha must be in [0, 1], got {tint_alpha}.") + + rgb = rendered[:, :3] + if base_images is None: + base_images = (rgb, rgb) + if len(base_images) != LAYER_CLASSES: + raise ValueError( + f"base_images must contain {LAYER_CLASSES} tensors, got {len(base_images)}." + ) + masks = geometry_debug[2:4] + edges = geometry_debug[4:6] + fill_colors = ( + rgb.new_tensor((0.10, 0.85, 0.95)), + rgb.new_tensor((1.00, 0.65, 0.10)), + ) + edge_colors = ( + rgb.new_tensor((0.00, 1.00, 1.00)), + rgb.new_tensor((1.00, 0.25, 0.85)), + ) + + overlays = [] + for base, mask, edge, fill_color, edge_color in zip( + base_images, masks, edges, fill_colors, edge_colors + ): + if base.shape != rgb.shape: + raise ValueError( + f"Geometry overlay base shape {tuple(base.shape)} does not match RGB {tuple(rgb.shape)}." + ) + fill_color = fill_color.view(1, 3, 1, 1) + edge_color = edge_color.view(1, 3, 1, 1) + tinted = base * (1.0 - tint_alpha) + fill_color * tint_alpha + overlay = torch.where(mask.unsqueeze(1), tinted, base) + overlay = torch.where(edge.unsqueeze(1), edge_color, overlay) + overlays.append(overlay) + return tuple(overlays) + + +def colorize_surface(labels, bg_color, reference): + """Colorize fixed renderer surface slots without assuming a fixed slot count.""" + N, H, W = labels.shape + bg = bg_tensor(bg_color, reference).expand(N, 3, H, W).clone() + valid = labels != IGNORE_INDEX + if not valid.any(): + return bg + surface = labels.clamp_min(0).to(dtype=reference.dtype) + colors = torch.stack( + [ + (surface.mul(47.0).remainder(191.0) + 32.0) / 255.0, + (surface.mul(83.0).remainder(191.0) + 32.0) / 255.0, + (surface.mul(131.0).remainder(191.0) + 32.0) / 255.0, + ], + dim=1, + ) + return torch.where(valid.unsqueeze(1), colors, bg) + + +def colorize_foreground(mask, bg_color, reference): + N, H, W = mask.shape + bg = bg_tensor(bg_color, reference).expand(N, 3, H, W) + fg = reference.new_ones(N, 3, H, W) + return torch.where(mask.unsqueeze(1), fg, bg) + + +def colorize_uv(uv, mask, bg_color): + N, _, H, W = uv.shape + bg = bg_tensor(bg_color, uv).expand(N, 3, H, W) + zeros = uv.new_zeros(N, 1, H, W) + uv_rgb = torch.cat([uv[:, 0:1], uv[:, 1:2], zeros], dim=1) + return torch.where(mask.unsqueeze(1), uv_rgb, bg) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_layout.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_layout.py new file mode 100644 index 0000000000000000000000000000000000000000..2230f8be36bc6a52a16ae9aad7d39c46d2507861 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_layout.py @@ -0,0 +1,171 @@ +"""Minecraft skin-atlas layout and RGBA helpers used by Dense UV Parser.""" + +import torch +import torchvision.transforms.functional as TF + + +UV_SIZE = 64 +PART_COUNT = 6 +FACE_COUNT = 6 +LAYER_COUNT = 2 + + +def minecraft_layer_rects(is_slim=False): + """Return inner face rectangles and their matching outer-layer offsets.""" + arm_width = 3 if is_slim else 4 + slim_shift = 1 if is_slim else 0 + parts = [ + ( + [ + ((8, 8), (8, 8)), + ((8, 8), (24, 8)), + ((8, 8), (16, 8)), + ((8, 8), (0, 8)), + ((8, 8), (8, 0)), + ((8, 8), (16, 0)), + ], + (32, 0), + ), + ( + [ + ((8, 12), (20, 20)), + ((8, 12), (32, 20)), + ((4, 12), (28, 20)), + ((4, 12), (16, 20)), + ((8, 4), (20, 16)), + ((8, 4), (28, 16)), + ], + (0, 16), + ), + ( + [ + ((arm_width, 12), (36, 52)), + ((arm_width, 12), (44 - slim_shift, 52)), + ((4, 12), (40 - slim_shift, 52)), + ((4, 12), (32, 52)), + ((arm_width, 4), (36, 48)), + ((arm_width, 4), (40 - slim_shift, 48)), + ], + (16, 0), + ), + ( + [ + ((arm_width, 12), (44, 20)), + ((arm_width, 12), (52 - slim_shift, 20)), + ((4, 12), (48 - slim_shift, 20)), + ((4, 12), (40, 20)), + ((arm_width, 4), (44, 16)), + ((arm_width, 4), (48 - slim_shift, 16)), + ], + (0, 16), + ), + ( + [ + ((4, 12), (20, 52)), + ((4, 12), (28, 52)), + ((4, 12), (24, 52)), + ((4, 12), (16, 52)), + ((4, 4), (20, 48)), + ((4, 4), (24, 48)), + ], + (-16, 0), + ), + ( + [ + ((4, 12), (4, 20)), + ((4, 12), (12, 20)), + ((4, 12), (8, 20)), + ((4, 12), (0, 20)), + ((4, 4), (4, 16)), + ((4, 4), (8, 16)), + ], + (0, 16), + ), + ] + + rects = [] + for faces, decor_offset in parts: + for (width, height), (inner_x, inner_y) in faces: + rects.append( + ( + inner_x, + inner_y, + width, + height, + decor_offset[0], + decor_offset[1], + ) + ) + return rects + + +def build_part_layer_masks(is_slim=False): + """Return exact inner/outer atlas masks grouped into six body parts.""" + inner = torch.zeros(PART_COUNT, 1, UV_SIZE, UV_SIZE, dtype=torch.float32) + outer = torch.zeros_like(inner) + for rect_index, (inner_x, inner_y, width, height, decor_dx, decor_dy) in enumerate( + minecraft_layer_rects(is_slim=is_slim) + ): + part = rect_index // FACE_COUNT + inner[part, :, inner_y : inner_y + height, inner_x : inner_x + width] = 1.0 + outer_x = inner_x + decor_dx + outer_y = inner_y + decor_dy + outer[part, :, outer_y : outer_y + height, outer_x : outer_x + width] = 1.0 + return inner, outer + + +def build_uv_masks(is_slim=False): + inner, outer = build_part_layer_masks(is_slim=is_slim) + return ( + inner.amax(dim=0), + outer.amax(dim=0), + ) + + +def finalize_minecraft_alpha( + tensor, + alpha_threshold=0.5, + enforce_base_alpha=True, + is_slim=False, +): + if tensor.shape[-3] != 4: + raise ValueError( + f"Expected RGBA tensor with 4 channels, got shape {tuple(tensor.shape)}." + ) + if not 0.0 <= alpha_threshold <= 1.0: + raise ValueError("alpha_threshold must be in [0, 1].") + + base_mask, decor_mask = build_uv_masks(is_slim=is_slim) + valid_mask = torch.maximum(base_mask, decor_mask) + base_mask = base_mask.to(device=tensor.device, dtype=tensor.dtype) + valid_mask = valid_mask.to(device=tensor.device, dtype=tensor.dtype) + + out = tensor.clone() + if tensor.dim() == 4: + base_mask = base_mask.unsqueeze(0) + valid_mask = valid_mask.unsqueeze(0) + alpha = (out[:, 3:4] > alpha_threshold).to(dtype=out.dtype) + if enforce_base_alpha: + alpha = torch.where(base_mask > 0, torch.ones_like(alpha), alpha) + out[:, :3] = out[:, :3] * valid_mask + out[:, 3:4] = alpha * valid_mask + return out + if tensor.dim() == 3: + alpha = (out[3:4] > alpha_threshold).to(dtype=out.dtype) + if enforce_base_alpha: + alpha = torch.where(base_mask > 0, torch.ones_like(alpha), alpha) + out[:3] = out[:3] * valid_mask + out[3:4] = alpha * valid_mask + return out + raise ValueError(f"Expected CHW or NCHW tensor, got shape {tuple(tensor.shape)}.") + + +def tensor_to_rgba_image(tensor): + tensor = tensor.detach().cpu().clamp(0.0, 1.0) + if tensor.shape[0] == 3: + tensor = torch.cat([tensor, torch.ones_like(tensor[:1])], dim=0) + return TF.to_pil_image(tensor) + + +def view_native_size(renderer, view): + return tuple(getattr(renderer, f"{view}_inner_mask").shape) diff --git a/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_topology.py b/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_topology.py new file mode 100644 index 0000000000000000000000000000000000000000..5903dc9402f786c3e9541866644145a9e20c2730 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/dense_uv_parser/uv_topology.py @@ -0,0 +1,259 @@ +"""Deterministic Minecraft UV topology needed by parser post-processing.""" + +from dataclasses import dataclass +from functools import lru_cache + +import torch + +from SkingToolkit.dense_uv_parser.uv_layout import ( + FACE_COUNT, + LAYER_COUNT, + PART_COUNT, + UV_SIZE, + minecraft_layer_rects, +) + + +INVALID_LAYER = LAYER_COUNT +INVALID_PART = PART_COUNT +INVALID_FACE = FACE_COUNT +MIRRORED_PART = (0, 1, 3, 2, 5, 4) +PART_CENTRES = ( + (0.0, 28.0, 0.0), + (0.0, 18.0, 0.0), + (6.0, 18.0, 0.0), + (-6.0, 18.0, 0.0), + (2.0, 6.0, 0.0), + (-2.0, 6.0, 0.0), +) + + +@dataclass(frozen=True) +class SimpleUVTopology: + valid: torch.Tensor + layer: torch.Tensor + part: torch.Tensor + face: torch.Tensor + local_uv: torch.Tensor + world_position: torch.Tensor + mirrored_texel: torch.Tensor + inner_fill_order: torch.Tensor + outer_flat_indices: torch.Tensor + outer_edge_index: torch.Tensor + + +def _axis_coordinate(index, size): + return -0.5 * float(size) + float(index) + 0.5 + + +def _surface_coordinate(face, u, v, width, height, depth): + x = _axis_coordinate(u, width) + y = _axis_coordinate(v, height) + z_from_front = 0.5 * depth - float(u) - 0.5 + z_from_back = -0.5 * depth + float(u) + 0.5 + if face == 0: + return (x, y, 0.5 * depth) + if face == 1: + return (-x, y, -0.5 * depth) + if face == 2: + return (0.5 * width, y, z_from_front) + if face == 3: + return (-0.5 * width, y, z_from_back) + if face == 4: + return (x, -0.5 * height, -0.5 * depth + float(v) + 0.5) + if face == 5: + return (x, 0.5 * height, -0.5 * depth + float(v) + 0.5) + raise ValueError(f"Unknown face index {face}.") + + +def _flat_index(x, y): + return int(y) * UV_SIZE + int(x) + + +def _inward_ring_key(cell, width, height): + u = cell["u"] + v = cell["v"] + ring = min(u, v, width - 1 - u, height - 1 - v) + right = width - 1 - ring + bottom = height - 1 - ring + if v == ring: + edge = 0 + offset = u - ring + elif u == right: + edge = 1 + offset = v - ring + elif v == bottom: + edge = 2 + offset = right - u + else: + edge = 3 + offset = bottom - v + return ring, edge, offset + + +@lru_cache(maxsize=1) +def build_simple_uv_topology(): + """Build Minecraft atlas metadata and the physical outer-surface graph.""" + valid = torch.zeros(UV_SIZE * UV_SIZE, dtype=torch.bool) + layer_map = torch.full((UV_SIZE * UV_SIZE,), INVALID_LAYER, dtype=torch.long) + part_map = torch.full((UV_SIZE * UV_SIZE,), INVALID_PART, dtype=torch.long) + face_map = torch.full((UV_SIZE * UV_SIZE,), INVALID_FACE, dtype=torch.long) + local_uv = torch.zeros(UV_SIZE * UV_SIZE, 2, dtype=torch.float32) + world_position = torch.zeros(UV_SIZE * UV_SIZE, 3, dtype=torch.float32) + cells_by_part = {} + + rects = minecraft_layer_rects(is_slim=False) + for layer in range(LAYER_COUNT): + for part in range(PART_COUNT): + part_rects = rects[part * FACE_COUNT : (part + 1) * FACE_COUNT] + width = part_rects[0][2] + height = part_rects[0][3] + depth = part_rects[2][2] + cells = [] + for face, ( + x, + y, + face_width, + face_height, + decor_dx, + decor_dy, + ) in enumerate(part_rects): + if layer == 1: + x += decor_dx + y += decor_dy + for v in range(face_height): + for u in range(face_width): + flat = _flat_index(x + u, y + v) + if valid[flat]: + raise ValueError( + f"Overlapping Minecraft atlas rectangles at {(x + u, y + v)}." + ) + valid[flat] = True + layer_map[flat] = layer + part_map[flat] = part + face_map[flat] = face + local_uv[flat, 0] = (u + 0.5) / max(face_width, 1) + local_uv[flat, 1] = (v + 0.5) / max(face_height, 1) + coordinate = _surface_coordinate( + face, u, v, width, height, depth + ) + if layer == 1: + expansion = 1.0 if part == 0 else 0.5 + coordinate = tuple( + value * (size + expansion) / size + for value, size in zip( + coordinate, (width, height, depth) + ) + ) + centre = PART_CENTRES[part] + world_position[flat] = torch.tensor( + [ + coordinate[0] + centre[0], + coordinate[1] + centre[1], + coordinate[2] + centre[2], + ], + dtype=torch.float32, + ) + cells.append( + { + "flat": flat, + "face": face, + "u": u, + "v": v, + } + ) + cells_by_part[(layer, part)] = cells + + mirrored = torch.arange(UV_SIZE * UV_SIZE, dtype=torch.long) + for layer in range(LAYER_COUNT): + for part in range(PART_COUNT): + targets = ( + valid & (layer_map == layer) & (part_map == part) + ).nonzero(as_tuple=False).flatten() + candidates = ( + valid + & (layer_map == layer) + & (part_map == MIRRORED_PART[part]) + ).nonzero(as_tuple=False).flatten() + reflected = world_position[targets].clone() + reflected[:, 0] = -reflected[:, 0] + nearest = torch.cdist( + reflected, world_position[candidates] + ).argmin(dim=1) + mirrored[targets] = candidates[nearest] + + inner_fill_order = [] + for part in range(PART_COUNT): + for face in range(FACE_COUNT): + cells = [ + cell + for cell in cells_by_part[(0, part)] + if cell["face"] == face + ] + width = max(cell["u"] for cell in cells) + 1 + height = max(cell["v"] for cell in cells) + 1 + if face in (2, 3): + cells.sort( + key=lambda cell: ( + cell["v"], + min(cell["u"], width - 1 - cell["u"]), + 0 if cell["u"] < width - 1 - cell["u"] else 1, + ) + ) + else: + cells.sort( + key=lambda cell: _inward_ring_key(cell, width, height) + ) + inner_fill_order.extend(cell["flat"] for cell in cells) + + outer_flat_indices = ( + valid & (layer_map == 1) + ).nonzero(as_tuple=False).flatten() + outer_edges = [] + # Atlas-neighbourhood alone is incorrect at cube seams: UV islands that + # touch in the PNG need not touch in 3D, while adjacent cube faces are + # usually separated in the atlas. World-space texel centres give us a + # deterministic face/seam-aware graph. The largest outer-layer step is + # 1.125 texels (limb/torso depth expansion), while a same-face diagonal is + # at least sqrt(2), leaving a useful gap for the threshold below. + for part in range(PART_COUNT): + part_nodes = outer_flat_indices[part_map[outer_flat_indices] == part] + positions = world_position[part_nodes] + distances = torch.cdist(positions, positions) + adjacent = (distances > 1e-6) & (distances <= 1.20) + source, target = adjacent.nonzero(as_tuple=True) + outer_edges.append(torch.stack([part_nodes[source], part_nodes[target]])) + outer_edge_index = torch.cat(outer_edges, dim=1) + + return SimpleUVTopology( + valid=valid.reshape(UV_SIZE, UV_SIZE), + layer=layer_map.reshape(UV_SIZE, UV_SIZE), + part=part_map.reshape(UV_SIZE, UV_SIZE), + face=face_map.reshape(UV_SIZE, UV_SIZE), + local_uv=local_uv.reshape(UV_SIZE, UV_SIZE, 2), + world_position=world_position.reshape(UV_SIZE, UV_SIZE, 3), + mirrored_texel=mirrored.reshape(UV_SIZE, UV_SIZE), + inner_fill_order=torch.tensor(inner_fill_order, dtype=torch.long), + outer_flat_indices=outer_flat_indices, + outer_edge_index=outer_edge_index, + ) + + +@lru_cache(maxsize=1) +def build_outer_uv_graph(): + """Return valid outer atlas indices and directed edges in node coordinates.""" + topology = build_simple_uv_topology() + flat_indices = topology.outer_flat_indices + flat_to_node = torch.full( + (UV_SIZE * UV_SIZE,), + -1, + dtype=torch.long, + ) + flat_to_node[flat_indices] = torch.arange( + flat_indices.numel(), + dtype=torch.long, + ) + edge_index = flat_to_node[topology.outer_edge_index] + if (edge_index < 0).any(): + raise ValueError("Outer topology contains an edge outside the outer atlas.") + return flat_indices, edge_index diff --git a/SKING_DDJ_v101c/runtime/renderer.py b/SKING_DDJ_v101c/runtime/renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..7fd1de35beda0a9f3ca2353a7d84632553509527 --- /dev/null +++ b/SKING_DDJ_v101c/runtime/renderer.py @@ -0,0 +1,295 @@ +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +class DifferentiableRenderer(nn.Module): + def __init__(self, mappings_dir=None, bg_color=(128/255, 128/255, 128/255)): + """ + Differentiable Minecraft Skin Renderer in PyTorch. + Args: + mappings_dir: Path to directory containing '*.pt' view mapping files. + If None, it tries to find the default mappings directory. + bg_color: RGB tuple for the render background. + """ + super().__init__() + if mappings_dir is None: + env_mappings_dir = os.environ.get("RENDERER_MAPPINGS_DIR") + if env_mappings_dir and os.path.exists(env_mappings_dir): + mappings_dir = env_mappings_dir + else: + try: + import sys + dmr_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "differentiable_minecraft_renderer")) + if dmr_dir not in sys.path: + sys.path.insert(0, dmr_dir) + from config import size as config_size + size_suffix = f"_{config_size[0]}x{config_size[1]}" + except Exception: + size_suffix = "" + + local_size_mappings = os.path.join(os.path.dirname(__file__), f"mappings{size_suffix}") + local_mappings = os.path.join(os.path.dirname(__file__), "mappings") + sibling_size_mappings = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "differentiable_minecraft_renderer", f"mappings{size_suffix}") + ) + sibling_mappings = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "differentiable_minecraft_renderer", "mappings") + ) + if os.path.exists(local_size_mappings): + mappings_dir = local_size_mappings + elif os.path.exists(local_mappings): + mappings_dir = local_mappings + elif os.path.exists(sibling_size_mappings): + mappings_dir = sibling_size_mappings + elif os.path.exists(sibling_mappings): + mappings_dir = sibling_mappings + else: + mappings_dir = "mappings" + + self.mappings_dir = mappings_dir + self.register_buffer("bg_color", torch.tensor(bg_color, dtype=torch.float32)) + + self.views = [] + + # Load and pre-process mappings dynamically + if os.path.exists(self.mappings_dir): + for file_name in sorted(os.listdir(self.mappings_dir)): + if file_name.endswith("_mapping.pt"): + view_name = file_name[:-11] # Remove "_mapping.pt" + self.views.append(view_name) + + mapping_path = os.path.join(self.mappings_dir, file_name) + # Load mapping coordinates dictionary + data = torch.load(mapping_path, map_location="cpu") + + # Extract and process inner layer + inner_uv = data["inner_uv_map"] # (H, W, 2) + inner_mask = data["inner_mask"] # (H, W) + + # Extract and process outer layer + outer_uv = data["outer_uv_map"] # (H, W, 2) + outer_mask = data["outer_mask"] # (H, W) + + # Normalize UV coordinates [0, 63] to [-1, 1] for F.grid_sample + inner_grid = torch.zeros_like(inner_uv) + inner_grid[..., 0] = (inner_uv[..., 0] / 63.0) * 2.0 - 1.0 + inner_grid[..., 1] = (inner_uv[..., 1] / 63.0) * 2.0 - 1.0 + + outer_grid = torch.zeros_like(outer_uv) + outer_grid[..., 0] = (outer_uv[..., 0] / 63.0) * 2.0 - 1.0 + outer_grid[..., 1] = (outer_uv[..., 1] / 63.0) * 2.0 - 1.0 + + # Register buffer so they are automatically moved to device with the module + self.register_buffer(f"{view_name}_inner_grid", inner_grid) + self.register_buffer(f"{view_name}_inner_mask", inner_mask) + self.register_buffer(f"{view_name}_outer_grid", outer_grid) + self.register_buffer(f"{view_name}_outer_mask", outer_mask) + + if "outer_uv_layers" in data and "outer_masks" in data: + outer_uv_layers = data["outer_uv_layers"].float() + outer_mask_layers = data["outer_masks"].float() + outer_grid_layers = torch.zeros_like(outer_uv_layers) + outer_grid_layers[..., 0] = (outer_uv_layers[..., 0] / 63.0) * 2.0 - 1.0 + outer_grid_layers[..., 1] = (outer_uv_layers[..., 1] / 63.0) * 2.0 - 1.0 + self.register_buffer(f"{view_name}_outer_grid_layers", outer_grid_layers) + self.register_buffer(f"{view_name}_outer_mask_layers", outer_mask_layers) + + if "composite_uv_layers" in data and "composite_masks" in data: + composite_uv_layers = data["composite_uv_layers"].float() + composite_mask_layers = data["composite_masks"].float() + composite_grid_layers = torch.zeros_like(composite_uv_layers) + composite_grid_layers[..., 0] = (composite_uv_layers[..., 0] / 63.0) * 2.0 - 1.0 + composite_grid_layers[..., 1] = (composite_uv_layers[..., 1] / 63.0) * 2.0 - 1.0 + self.register_buffer(f"{view_name}_composite_grid_layers", composite_grid_layers) + self.register_buffer(f"{view_name}_composite_mask_layers", composite_mask_layers) + if "composite_is_decor_layers" in data: + self.register_buffer( + f"{view_name}_composite_is_decor_layers", + data["composite_is_decor_layers"].bool(), + ) + + if ( + "geometry_uv_layers" in data + and "geometry_masks" in data + and ("geometry_sort_indices" in data or "geometry_depth_layers" in data) + ): + geometry_uv_layers = data["geometry_uv_layers"].float() + geometry_mask_layers = data["geometry_masks"].float() + geometry_grid_layers = torch.zeros_like(geometry_uv_layers) + geometry_grid_layers[..., 0] = (geometry_uv_layers[..., 0] / 63.0) * 2.0 - 1.0 + geometry_grid_layers[..., 1] = (geometry_uv_layers[..., 1] / 63.0) * 2.0 - 1.0 + if "geometry_sort_indices" in data: + geometry_sort_indices = data["geometry_sort_indices"].long() + else: + geometry_depth_layers = data["geometry_depth_layers"].float() + geometry_sort_indices = torch.argsort(geometry_depth_layers, dim=0) + self.register_buffer(f"{view_name}_geometry_grid_layers", geometry_grid_layers) + self.register_buffer(f"{view_name}_geometry_mask_layers", geometry_mask_layers) + self.register_buffer(f"{view_name}_geometry_sort_indices", geometry_sort_indices) + else: + print(f"WARNING: mappings_dir '{self.mappings_dir}' does not exist yet. Views list is empty.") + + def forward_view(self, skins, view_name): + """ + Render a single view for a batch of skins. + Args: + skins: PyTorch tensor of shape (B, 4, 64, 64) with values in range [0, 1]. + The channel order is RGBA. + view_name: One of registered view names (e.g. "front", "back", "static_front"). + Returns: + rendered: PyTorch tensor of shape (B, 4, H, W) containing RGBA render. + """ + B, C, H_in, W_in = skins.shape + assert C == 4, "Skins must have 4 channels (RGBA)" + assert H_in == 64 and W_in == 64, "Skins must be 64x64" + + # Retrieve buffers and cast to match the input dtype (e.g. bfloat16 or float32) + dtype = skins.dtype + inner_grid = getattr(self, f"{view_name}_inner_grid").unsqueeze(0).expand(B, -1, -1, -1).to(dtype=dtype) + inner_mask = getattr(self, f"{view_name}_inner_mask").unsqueeze(0).unsqueeze(1).expand(B, -1, -1, -1).to(dtype=dtype) # (B, 1, H, W) + + outer_grid = getattr(self, f"{view_name}_outer_grid").unsqueeze(0).expand(B, -1, -1, -1).to(dtype=dtype) + outer_mask = getattr(self, f"{view_name}_outer_mask").unsqueeze(0).unsqueeze(1).expand(B, -1, -1, -1).to(dtype=dtype) # (B, 1, H, W) + + # 1. Sample inner layer using bilinear interpolation + inner_sampled = F.grid_sample(skins, inner_grid, mode='bilinear', padding_mode='zeros', align_corners=True) + inner_sampled = inner_sampled * inner_mask + + # 2. Sample outer layer using bilinear interpolation + outer_sampled = F.grid_sample(skins, outer_grid, mode='bilinear', padding_mode='zeros', align_corners=True) + outer_sampled = outer_sampled * outer_mask + + # 3. Alpha blend outer over inner + inner_rgb = inner_sampled[:, :3, :, :] + inner_alpha = inner_sampled[:, 3:4, :, :] + + outer_rgb = outer_sampled[:, :3, :, :] + outer_alpha = outer_sampled[:, 3:4, :, :] + + # Prepare background color broadcasted to match dimensions + bg = self.bg_color.view(1, 3, 1, 1).expand(B, -1, inner_rgb.shape[2], inner_rgb.shape[3]).to(dtype=dtype) + + # Composite inner layer over background + inner_composite = inner_alpha * inner_rgb + (1.0 - inner_alpha) * bg + + # Composite outer layer over inner composite + final_rgb = outer_alpha * outer_rgb + (1.0 - outer_alpha) * inner_composite + + # Calculate final alpha mask (for background/foreground separation) + final_alpha = outer_alpha + (1.0 - outer_alpha) * inner_alpha + + composite_grid_name = f"{view_name}_composite_grid_layers" + composite_decor_name = f"{view_name}_composite_is_decor_layers" + if hasattr(self, composite_grid_name) and hasattr(self, composite_decor_name): + composite_grid_layers = getattr(self, composite_grid_name).to(dtype=dtype) + composite_mask_layers = getattr(self, f"{view_name}_composite_mask_layers").to(dtype=dtype) + composite_is_decor_layers = getattr(self, composite_decor_name) + num_layers = composite_grid_layers.shape[0] + out_h, out_w = composite_grid_layers.shape[1:3] + + composite_grid = composite_grid_layers.unsqueeze(0).expand(B, -1, -1, -1, -1) + composite_grid = composite_grid.reshape(B * num_layers, out_h, out_w, 2) + skin_layers = skins.unsqueeze(1).expand(-1, num_layers, -1, -1, -1) + skin_layers = skin_layers.reshape(B * num_layers, C, H_in, W_in) + composite_sampled = F.grid_sample( + skin_layers, + composite_grid, + mode='bilinear', + padding_mode='zeros', + align_corners=True, + ) + composite_sampled = composite_sampled.reshape(B, num_layers, C, out_h, out_w) + composite_mask = composite_mask_layers.unsqueeze(0).unsqueeze(2).expand(B, -1, -1, -1, -1) + composite_sampled = composite_sampled * composite_mask + + composite_render_rgb = bg + composite_render_alpha = torch.zeros_like(final_alpha) + for layer_index in range(num_layers - 1, -1, -1): + layer = composite_sampled[:, layer_index] + layer_rgb = layer[:, :3] + layer_alpha = layer[:, 3:4] + composite_render_rgb = layer_alpha * layer_rgb + (1.0 - layer_alpha) * composite_render_rgb + composite_render_alpha = layer_alpha + (1.0 - layer_alpha) * composite_render_alpha + + layer_alpha = composite_sampled[:, :, 3:4] + visible = (layer_alpha > 1e-4) & (composite_mask > 0.5) + layer_indices = torch.arange(num_layers, device=skins.device).view(1, num_layers, 1, 1, 1) + fallback_indices = torch.full(visible.shape, num_layers, device=skins.device, dtype=torch.long) + visible_indices = torch.where(visible, layer_indices, fallback_indices) + first_visible_index = visible_indices.amin(dim=1) + has_visible = first_visible_index < num_layers + + decor_layers = composite_is_decor_layers.unsqueeze(0).unsqueeze(2).expand(B, -1, -1, -1, -1) + first_layer = layer_indices == first_visible_index.unsqueeze(1) + first_visible_is_decor = (visible & first_layer & decor_layers).any(dim=1) + trust_composite = has_visible & (first_visible_is_decor | (first_visible_index <= 1)) + + final_rgb = torch.where(trust_composite, composite_render_rgb, final_rgb) + final_alpha = torch.where(trust_composite, composite_render_alpha, final_alpha) + + geometry_grid_name = f"{view_name}_geometry_grid_layers" + if hasattr(self, geometry_grid_name) and num_layers > 0: + geometry_grid_layers = getattr(self, geometry_grid_name).to(dtype=dtype) + geometry_mask_layers = getattr(self, f"{view_name}_geometry_mask_layers").to(dtype=dtype) + geometry_sort_indices = getattr(self, f"{view_name}_geometry_sort_indices").to(device=skins.device) + geometry_layers = geometry_grid_layers.shape[0] + if geometry_layers > 0: + geometry_grid = geometry_grid_layers.unsqueeze(0).expand(B, -1, -1, -1, -1) + geometry_grid = geometry_grid.reshape(B * geometry_layers, out_h, out_w, 2) + geometry_skin_layers = skins.unsqueeze(1).expand(-1, geometry_layers, -1, -1, -1) + geometry_skin_layers = geometry_skin_layers.reshape(B * geometry_layers, C, H_in, W_in) + geometry_sampled = F.grid_sample( + geometry_skin_layers, + geometry_grid, + mode='bilinear', + padding_mode='zeros', + align_corners=True, + ) + geometry_sampled = geometry_sampled.reshape(B, geometry_layers, C, out_h, out_w) + geometry_mask = geometry_mask_layers.unsqueeze(0).unsqueeze(2).expand( + B, -1, -1, -1, -1 + ) + geometry_sampled = geometry_sampled * geometry_mask + geometry_order = geometry_sort_indices.unsqueeze(0).unsqueeze(2).expand( + B, -1, C, -1, -1 + ) + geometry_sampled = torch.gather(geometry_sampled, dim=1, index=geometry_order) + + geometry_render_rgb = bg + geometry_render_alpha = torch.zeros_like(final_alpha) + for layer_index in range(geometry_layers): + layer = geometry_sampled[:, layer_index] + layer_rgb = layer[:, :3] + layer_alpha = layer[:, 3:4] + geometry_render_rgb = ( + layer_alpha * layer_rgb + (1.0 - layer_alpha) * geometry_render_rgb + ) + geometry_render_alpha = layer_alpha + (1.0 - layer_alpha) * geometry_render_alpha + + front_decor = composite_is_decor_layers[0].unsqueeze(0).unsqueeze(1) + front_mask = (composite_mask_layers[0] > 0.5).unsqueeze(0).unsqueeze(1) + front_alpha = composite_sampled[:, 0, 3:4] + geometry_fallback = ( + front_mask + & front_decor + & (front_alpha <= 1e-4) + & (geometry_render_alpha > 1e-4) + ) + final_rgb = torch.where(geometry_fallback, geometry_render_rgb, final_rgb) + final_alpha = torch.where(geometry_fallback, geometry_render_alpha, final_alpha) + + return torch.cat([final_rgb, final_alpha], dim=1) + + def forward(self, skins): + """ + Renders all loaded views for a batch of skins. + Args: + skins: PyTorch tensor of shape (B, 4, 64, 64) with values in range [0, 1]. + Returns: + dict of rendered views. Each view has shape (B, 4, H, W). + """ + results = {} + for view_name in self.views: + results[view_name] = self.forward_view(skins, view_name) + return results diff --git a/SKING_DDJ_v101c/tested_environments.json b/SKING_DDJ_v101c/tested_environments.json new file mode 100644 index 0000000000000000000000000000000000000000..55211c6ef4133d3cd6a99d2035ad26b5bc125277 --- /dev/null +++ b/SKING_DDJ_v101c/tested_environments.json @@ -0,0 +1,34 @@ +{ + "parser": { + "python": "3.12", + "torch": "2.8.0+cu128", + "torchvision": "0.23.0+cu128", + "numpy": "2.4.1", + "Pillow": "12.1.0", + "transformers": "5.0.0", + "huggingface-hub": "1.3.4", + "timm": "1.0.24", + "einops": "0.8.1", + "sentencepiece": "0.2.1", + "safetensors": "0.7.0", + "accelerate": "1.12.0", + "scipy": "1.18.0", + "opencv-python": "5.0.0.93" + }, + "foreground": { + "python": "3.12", + "torch": "2.8.0+cu128", + "torchvision": "0.23.0+cu128", + "numpy": "2.4.1", + "Pillow": "12.1.0", + "transformers": "5.0.0", + "huggingface-hub": "1.3.4", + "timm": "1.0.24", + "einops": "0.8.1", + "sentencepiece": "0.2.1", + "safetensors": "0.7.0", + "accelerate": "1.12.0", + "scipy": "1.17.0", + "kornia": "0.8.2" + } +}