Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,22 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.ipynb_checkpoints/
+.vscode/
+
+# Weights in output/scratch directories (tracked files in model subfolders use LFS via .gitattributes)
+eval_output/**/*.safetensors
+eval_output/**/*.pth
+eval_output/**/*.bin
+eval_output/**/*.ckpt
+
+# Output directories
+eval_output/
+infer_output/
+visualizations/
+
+# Virtual environment
+.venv/
+
+# External Repositories
+dinov3/

Firefly_RGB/config.json

@@ -0,0 +1,47 @@
+{
+  "architectures": ["FireflyForSemanticSegmentationRGB"],
+  "model_type": "firefly-rgb",
+  "num_labels": 15,
+  "backbone": "dinov3_vitb16",
+  "image_size": 640,
+  "embedding_dim": 256,
+  "backbone_embed_dim": 768,
+  "patch_size": 16,
+  "dropout_ratio": 0.1,
+  "semantic_loss_ignore_index": 255,
+  "id2label": {
+    "0": "road",
+    "1": "walkway",
+    "2": "dirt",
+    "3": "gravel",
+    "4": "grass",
+    "5": "vegetation",
+    "6": "tree",
+    "7": "ground_obstacles",
+    "8": "vehicle",
+    "9": "water",
+    "10": "building",
+    "11": "roof",
+    "12": "parking_lot",
+    "13": "constructions",
+    "14": "truck"
+  },
+  "label2id": {
+    "road": 0,
+    "walkway": 1,
+    "dirt": 2,
+    "gravel": 3,
+    "grass": 4,
+    "vegetation": 5,
+    "tree": 6,
+    "ground_obstacles": 7,
+    "vehicle": 8,
+    "water": 9,
+    "building": 10,
+    "roof": 11,
+    "parking_lot": 12,
+    "constructions": 13,
+    "truck": 14
+  },
+  "transformers_version": "4.37.2"
+}

Firefly_RGB/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71d74d0fb68508e515c97a09081ca6e613f9bef4dfae8fc00f0b6f3fcdf2be41
+size 352263468

Firefly_RGB/preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "processor_class": "SegformerImageProcessor",
+  "do_normalize": true,
+  "do_reduce_labels": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "SegformerImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 0,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 640,
+    "width": 640
+  }
+}

Firefly_Thermal/config.json

@@ -0,0 +1,50 @@
+{
+  "architectures": ["FireflyForSemanticSegmentationThermal"],
+  "model_type": "firefly-thermal",
+  "num_labels": 15,
+  "backbone": "dinov3_vitb16",
+  "image_size": 640,
+  "embedding_dim": 256,
+  "backbone_embed_dim": 768,
+  "patch_size": 16,
+  "num_layers": 12,
+  "rein_token_length": 100,
+  "feature_layers": [2, 5, 8, 11],
+  "dropout_ratio": 0.1,
+  "semantic_loss_ignore_index": 255,
+  "id2label": {
+    "0": "road",
+    "1": "walkway",
+    "2": "dirt",
+    "3": "gravel",
+    "4": "grass",
+    "5": "vegetation",
+    "6": "tree",
+    "7": "ground_obstacles",
+    "8": "vehicle",
+    "9": "water",
+    "10": "building",
+    "11": "roof",
+    "12": "parking_lot",
+    "13": "constructions",
+    "14": "truck"
+  },
+  "label2id": {
+    "road": 0,
+    "walkway": 1,
+    "dirt": 2,
+    "gravel": 3,
+    "grass": 4,
+    "vegetation": 5,
+    "tree": 6,
+    "ground_obstacles": 7,
+    "vehicle": 8,
+    "water": 9,
+    "building": 10,
+    "roof": 11,
+    "parking_lot": 12,
+    "constructions": 13,
+    "truck": 14
+  },
+  "transformers_version": "4.37.2"
+}

Firefly_Thermal/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc3e5f8c6e3426dd8379ae7ca476138e76e76adfbd890e23b77ce49077eb69f0
+size 362250440

Firefly_Thermal/preprocessor_config.json

@@ -0,0 +1,24 @@
+{
+  "processor_class": "SegformerImageProcessor",
+  "do_normalize": true,
+  "do_reduce_labels": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.3990,
+    0.3990,
+    0.3990
+  ],
+  "image_processor_type": "SegformerImageProcessor",
+  "image_std": [
+    0.2009,
+    0.2009,
+    0.2009
+  ],
+  "resample": 0,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 640,
+    "width": 640
+  }
+}

LICENSE

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

README.md

@@ -0,0 +1,153 @@
+---
+license: cc-by-nc-sa-4.0
+library_name: transformers
+pipeline_tag: image-segmentation
+tags:
+  - semantic-segmentation
+  - drone
+  - rgb
+  - thermal
+  - infrared
+  - dinov3
+  - aerial
+datasets:
+  - markus-42/SegFly
+model-index:
+  - name: Firefly-RGB
+    results: []
+  - name: Firefly-Thermal
+    results: []
+---
+
+# Introduction
+
+<div align="center">
+
+[![GitHub Documentation](https://img.shields.io/badge/GitHub-Documentation-green?logo=github&logoColor=white&labelColor=555)](https://github.com/markus-42/SegFly)&nbsp;&nbsp;
+[![Project](https://img.shields.io/badge/Project-Page-blue.svg)](https://markus-42.github.io/publications/2026/segfly/)&nbsp;&nbsp;
+[![arXiv](https://img.shields.io/badge/arXiv-Paper-red.svg)](https://arxiv.org/abs/2603.17920)&nbsp;&nbsp;
+[![SegFly Dataset](https://img.shields.io/badge/HuggingFace-SegFly%20Dataset%20-e58f00?logo=huggingface&logoColor=white&labelColor=555)](https://huggingface.co/datasets/markus-42/SegFly)
+
+</div>
+
+Following the acceptance of [SegFly](https://markus-42.github.io/publications/2026/segfly/) at the ECCV 2026 computer vision conference, we not only release the [SegFly dataset](https://huggingface.co/datasets/markus-42/SegFly) on HuggingFace, but also our `Firefly` semantic segmentation model, specialized for RGB and thermal imagery from aerial perspectives. Firefly was trained on our SegFly dataset.
+
+#### Key Features:
+- **Aerial-specialized**: Fine-tuned on diverse aerial imagery from urban, industrial, and rural environments.
+- **Multi-altitude performance**: Trained on data from 50m, 40m, and 30m altitudes.
+- **Seasonal robustness**: Covers data across all seasons for improved generalization.
+- **Semantic Categories**: Both models segment 15 semantic classes: road, walkway, dirt, gravel, grass, vegetation, tree, ground obstacles, vehicle, water, building, roof, parking lot, constructions, and truck.
+
+
+# Checkpoints
+
+Two pretrained checkpoints are provided:
+
+| Variant | Subfolder | Modality | Architecture |
+| :--- | :--- | :--- | :--- |
+| Firefly-RGB | `Firefly_RGB/` | RGB drone imagery | DINOv3 ViT-B/16 + MLP head |
+| Firefly-Thermal | `Firefly_Thermal/` | Thermal drone imagery | DINOv3 ViT-B/16 + Rein adapter + MLP head |
+
+
+# Setup
+
+```bash
+git clone https://huggingface.co/markus-42/SegFly-Firefly
+cd SegFly-Firefly
+
+# clone DINOv3 backbone inside the repository
+git clone https://github.com/facebookresearch/dinov3.git
+
+# create virtual environment and install dependencies
+uv venv && uv pip install -r requirements.txt
+
+# activate virtual environment
+source .venv/bin/activate
+```
+
+# Inference
+
+Run single-image inference with `infer.py`. Weights are auto-detected from the modality:
+
+```bash
+# RGB model
+python infer.py --image example.jpg --modality rgb
+
+# Thermal model
+python infer.py --image example_thermal.jpg --modality thermal
+```
+
+The colorized segmentation map is saved to `./infer_output/` by default.
+
+| Argument | Type | Default | Description |
+| :--- | :--- | :--- | :--- |
+| `--image` | `str` | *(required)* | Path to the input image. |
+| `--modality` | `str` | `rgb` | Model to use: `rgb` or `thermal`. |
+| `--output` | `str` | `./infer_output` | Directory for the colorized segmentation output. |
+| `--weights_path` | `str` | `""` | Path to model weights. Auto-detected from modality if not set. |
+| `--dinov3_repo_dir` | `str` | `./dinov3` | Path to the local DINOv3 repository. |
+| `--image_size` | `int` | `640` | Resolution at which the image is fed to the model. |
+
+# Batch Evaluation
+
+`eval.py` computes segmentation metrics (mIoU, Frequency-Weighted IoU, Pixel Accuracy) against ground truth masks for a set of images.
+
+```bash
+# RGB model
+python eval.py --data_dir /path/to/dataset --modality rgb
+
+# Thermal model
+python eval.py --data_dir /path/to/dataset --modality thermal
+```
+
+`--data_dir` can point to any directory. The loader walks it recursively and collects images and masks by subfolder name — no specific top-level structure is required:
+
+```
+data_dir/
+└── <any nesting>/
+    ├── src/           ← RGB images (.png / .jpg / .jpeg)       [--modality rgb]
+    ├── thermal_src/   ← thermal images (.png / .jpg / .jpeg)   [--modality thermal]
+    └── gt/            ← grayscale segmentation masks (.png, raw SegFly class IDs)
+```
+
+Images and masks are matched by alphabetical sort order, so filenames must correspond 1-to-1.
+
+Append `--visualize` to save colorized segmentation maps to `<output_dir>/visualizations`.
+
+| Argument | Type | Default | Description |
+| :--- | :--- | :--- | :--- |
+| `--data_dir` | `str` | `./data` | Root directory of the dataset. Walked recursively; images are collected from `src/` or `thermal_src/` subfolders and masks from `gt/` subfolders. |
+| `--weights_path` | `str` | `""` | Path to the saved weights. Auto-detected from modality if not set. |
+| `--class_dict_path` | `str` | `./classes_segfly.csv` | Path to the CSV file defining the dataset's class mapping. |
+| `--modality` | `str` | `rgb` | Modality of the dataset: `rgb` or `thermal`. |
+| `--output_dir` | `str` | `./eval_output` | Directory to save `per_image_iou.csv`, `per_class_scores.txt`, and visualizations. |
+| `--visualize` | `flag` | `False` | If provided, generates side-by-side visualizations of Image, Ground Truth, and Prediction. |
+| `--image_size` | `int` | `640` | Input resolution for evaluation. |
+| `--dinov3_repo_dir` | `str` | `./dinov3` | Path to the local DINOv3 repository. |
+
+
+# Reference
+If our work was helpful to you, we would appreciate citing our papers and giving the repository a like ❤️
+
+```bibtex
+@inproceedings{gross2026segfly,
+    title={{SegFly: A Dataset and 2D-3D-2D Paradigm for Aerial RGB-Thermal Semantic Segmentation at Scale}}, 
+    author={Markus Gross and Sai Bharadhwaj Matha and Rui Song and Viswanathan Muthuveerappan and Conrad Christoph and Julius Huber and Daniel Cremers},
+    booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
+    year={2026},
+}
+```
+
+Since SegFly is based on the [OccuFly dataset](https://markus-42.github.io/publications/2026/occufly/), consider citing this work as well:
+
+```bibtex
+@inproceedings{gross2026occufly,
+    title={{OccuFly: A 3D Vision Benchmark for Semantic Scene Completion from the Aerial Perspective}}, 
+    author={Markus Gross and Sai B. Matha and Aya Fahmy and Rui Song and Daniel Cremers and Henri Meess},
+    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    year={2026},
+}
+```
+
+# License
+This work is licensed under the [CC BY-NC-SA 4.0 license](https://creativecommons.org/licenses/by-nc-sa/4.0/). See the LICENSE file for the full legal terms.

classes_segfly.csv

@@ -0,0 +1,17 @@
+id,name
+0,"empty"
+1,"road"
+2,"walkway"
+3,"dirt"
+4,"gravel"
+6,"grass"
+7,"vegetation"
+8,"tree"
+9,"ground_obstacles"
+13,"vehicle"
+14,"water"
+16,"building"
+17,"roof"
+33,"parking_lot"
+34,"constructions"
+36,"truck"

eval.py

@@ -0,0 +1,447 @@
+import os
+import argparse
+import evaluate
+import numpy as np
+import torch
+import csv
+import pandas as pd
+from tqdm.auto import tqdm
+import cv2
+import re
+
+from lib.utils_segfly import InferenceDataset, InferenceDatasetThermal, Timing, ID2COLOR, mask2label
+from safetensors.torch import load_file
+from transformers import AutoImageProcessor
+from lib.firefly_rgb import FireflyForSemanticSegmentationRGB, FireflyConfigRGB
+from lib.firefly_thermal import FireflyForSemanticSegmentationThermal, FireflyConfigThermal
+
+timing = Timing()
+
+
+@timing
+def segment_image(image, _model, _device, target_size, image_processor):
+    with torch.no_grad(): 
+        pixel_values = image_processor(image, return_tensors="pt").pixel_values.to(_device, dtype=torch.float32)
+        outputs = _model(pixel_values)
+        
+        if hasattr(outputs, "logits"):
+            outputs = outputs.logits
+            
+        upsampled_logits = torch.nn.functional.interpolate(
+            outputs.float(), size=target_size, mode="bilinear", align_corners=False
+        )
+        return upsampled_logits.argmax(dim=1).detach().cpu().numpy()
+
+def run_inference(image, _model, _device, target_size, image_processor):
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    
+    predicted = segment_image(image, _model, _device, target_size, image_processor)
+        
+    mem_used = torch.cuda.max_memory_allocated() / 1024**2
+    return predicted, mem_used
+
+def get_image_path_info(dataset, idx):
+    path = ""
+    for attr in ['images', 'image_paths', 'filepaths', 'samples', 'img_files', 'data']:
+        if hasattr(dataset, attr):
+            val = getattr(dataset, attr)
+            if isinstance(val, (list, tuple, np.ndarray)) and len(val) > idx:
+                path = str(val[idx])
+                break
+    
+    subfolder_name = f"pred_img_{idx:04d}"
+    if path:
+        scene_match = re.search(r'(scene_\d+)', path, re.IGNORECASE)
+        alt_match = re.search(r'(\d+m)', path, re.IGNORECASE)
+        
+        if scene_match or alt_match:
+            scene_str = scene_match.group(1) if scene_match else "scene_unk"
+            alt_str = alt_match.group(1) if alt_match else "unk_m"
+            subfolder_name = f"pred_{scene_str}_{alt_str}_{idx:04d}"
+            
+    actual_path = path
+    if path and not os.path.exists(path):
+        if hasattr(dataset, 'root_dir') and os.path.exists(os.path.join(dataset.root_dir, path)):
+            actual_path = os.path.join(dataset.root_dir, path)
+        elif hasattr(dataset, '_root_dir') and os.path.exists(os.path.join(dataset._root_dir, path)):
+            actual_path = os.path.join(dataset._root_dir, path)
+
+    return subfolder_name, actual_path
+
+def generate_metrics_for_dataset(loader, _model, _device, _metric, config, image_processor, visualize=False):
+    mem_usage = []
+    global_gt_counts = np.zeros(config["num_classes"], dtype=np.int64)
+    ignore_label = config.get("ignore_label", 255)
+    
+    if visualize:
+        vis_dir = os.path.join(config.get("output_dir", "./"), "visualizations")
+        os.makedirs(vis_dir, exist_ok=True)
+        palette = ID2COLOR
+
+    csv_path = os.path.join(config.get("output_dir", "./"), "per_image_iou.csv")
+    
+    print("Running inference and accumulating batches...")
+    pbar = tqdm(total=len(loader)) 
+    
+    with open(csv_path, mode="w", newline="", encoding="utf-8") as csv_file:
+        csv_writer = csv.writer(csv_file)
+        csv_writer.writerow(["Image_Index", "mIoU"])
+        
+        for idx, (images, masks) in enumerate(loader):
+            image = images.squeeze(0) 
+            mask = masks.squeeze(0)
+            
+            if len(mask.shape) == 2:
+                target_h, target_w = mask.shape
+            else:
+                target_h, target_w = mask.shape[-2:]
+            
+            target_size = (target_h, target_w)
+
+            prediction, mem = run_inference(image, _model, _device, target_size, image_processor)
+            
+            if isinstance(mask, torch.Tensor):
+                mask_np = mask.detach().cpu().numpy().astype(int)
+            else:
+                mask_np = np.array(mask).astype(int)
+            
+            if getattr(image_processor, 'do_reduce_labels', False):
+                mask_np = mask_np.copy()
+                mask_np[mask_np == 0] = 255
+                mask_np = mask_np - 1
+                mask_np[mask_np == 254] = 255
+
+            valid_pixels = mask_np[mask_np != ignore_label]
+            valid_pixels = valid_pixels[valid_pixels >= 0] 
+            counts = np.bincount(valid_pixels.flatten(), minlength=config["num_classes"])
+            global_gt_counts += counts[:config["num_classes"]]
+            
+            pred_np = prediction.squeeze()
+            if isinstance(pred_np, torch.Tensor):
+                pred_np = pred_np.cpu().numpy()
+            
+            valid = (mask_np != ignore_label) & (mask_np >= 0)
+            pred_valid = pred_np[valid]
+            mask_valid = mask_np[valid]
+            
+            if len(mask_valid) > 0:
+                classes_in_img = np.unique(np.concatenate([mask_valid, pred_valid]))
+            else:
+                classes_in_img = []
+                
+            image_ious = []
+            for c in classes_in_img:
+                intersection = np.sum((pred_valid == c) & (mask_valid == c))
+                union = np.sum((pred_valid == c) | (mask_valid == c))
+                if union > 0:
+                    image_ious.append(intersection / union)
+                else:
+                    image_ious.append(0.0)
+            
+            image_miou = np.mean(image_ious) if len(image_ious) > 0 else 0.0
+            csv_writer.writerow([idx, f"{image_miou:.4f}"])
+            csv_file.flush()
+            
+            if visualize:
+                subfolder_name, actual_image_path = get_image_path_info(loader.dataset, idx)
+                
+                true_img_bgr = None
+                if actual_image_path and os.path.exists(actual_image_path):
+                    true_img_bgr = cv2.imread(actual_image_path)
+                
+                if true_img_bgr is not None:
+                    orig_h, orig_w = true_img_bgr.shape[:2]
+                    img_bgr = true_img_bgr
+                else:
+                    print(f"\n[Warning] Could not load raw image from disk for idx {idx}. Falling back to tensor size.")
+                    if isinstance(image, torch.Tensor):
+                        img_vis = image.cpu().numpy()
+                    else:
+                        img_vis = np.array(image)
+                    
+                    if img_vis.ndim == 3:
+                        if img_vis.shape[0] in [1, 3]:       
+                            img_vis = np.transpose(img_vis, (1, 2, 0))
+                        elif img_vis.shape[1] in [1, 3]:     
+                            img_vis = np.transpose(img_vis, (0, 2, 1))
+                            
+                    if img_vis.dtype.kind == 'f' and img_vis.max() <= 1.0:
+                        img_vis = img_vis * 255.0
+                    
+                    img_vis = np.ascontiguousarray(img_vis).astype(np.uint8)
+                    
+                    if img_vis.ndim == 2:
+                        img_vis = cv2.cvtColor(img_vis, cv2.COLOR_GRAY2RGB)
+                    elif img_vis.ndim == 3 and img_vis.shape[2] == 1:
+                        img_vis = cv2.cvtColor(img_vis[:, :, 0], cv2.COLOR_GRAY2RGB)
+                    elif img_vis.ndim == 3 and img_vis.shape[2] > 3:
+                        img_vis = img_vis[:, :, :3]
+                        
+                    orig_h, orig_w = img_vis.shape[:2]
+                    img_bgr = cv2.cvtColor(img_vis, cv2.COLOR_RGB2BGR)
+                
+                gt_color = mask2label(mask_np, palette)
+                pred_color = mask2label(pred_np, palette)
+
+                gt_color[mask_np == ignore_label] = [0, 0, 0]
+
+                try:
+                    if gt_color.shape[:2] != (orig_h, orig_w):
+                        gt_color = cv2.resize(gt_color, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+                    if pred_color.shape[:2] != (orig_h, orig_w):
+                        pred_color = cv2.resize(pred_color, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+                except cv2.error as e:
+                    print(f"\n[Error] OpenCV failed to resize masks. Original mask shape: {gt_color.shape[:2]}, Target image size: ({orig_w}, {orig_h})")
+                    raise e
+
+                instance_dir = os.path.join(vis_dir, subfolder_name)
+                os.makedirs(instance_dir, exist_ok=True)
+
+                gt_bgr = cv2.cvtColor(gt_color, cv2.COLOR_RGB2BGR)
+                pred_bgr = cv2.cvtColor(pred_color, cv2.COLOR_RGB2BGR)
+
+                cv2.imwrite(os.path.join(instance_dir, "image.png"), img_bgr)
+                cv2.imwrite(os.path.join(instance_dir, "gt.png"), gt_bgr)
+                cv2.imwrite(os.path.join(instance_dir, "pred.png"), pred_bgr)
+
+            if len(prediction.shape) == 2:
+                prediction = prediction[None, :, :]
+            if len(mask.shape) == 2:
+                mask = mask[None, :, :]
+
+            _metric.add_batch(
+                predictions=prediction,
+                references=mask
+            )
+            
+            mem_usage.append(mem)
+            pbar.update()
+            
+    pbar.close()
+    
+    print("Computing global metrics...")
+    results = _metric.compute(
+        num_labels=config["num_classes"],
+        ignore_index=config.get("ignore_label"),
+        reduce_labels=image_processor.do_reduce_labels,
+    )
+    
+    return results, mem_usage, global_gt_counts
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Evaluation script for Firefly RGB / Thermal models")
+    parser.add_argument("--data_dir", type=str, default="./data", help="Path to the root directory of the dataset.")
+    parser.add_argument("--weights_path", type=str, default="", help="Path to the .safetensors or .bin or .pth trained weights file.")
+    parser.add_argument("--class_dict_path", type=str, default="./classes_segfly.csv", help="Path to the class dictionary CSV file.")
+    parser.add_argument("--modality", type=str, default="rgb", choices=["rgb", "thermal"], help="Modality of the dataset (rgb or thermal).")
+    parser.add_argument("--output_dir", type=str, default="./eval_output", help="Directory where evaluation results and visualizations will be saved.")
+    parser.add_argument("--visualize", action="store_true", help="Generate and save side-by-side visualizations of image, ground truth, and prediction.")
+    parser.add_argument("--image_size", type=int, default=640, help="Image size for inference.")
+    parser.add_argument("--dinov3_repo_dir", type=str, default="./dinov3", help="Path to the local dinov3 repository.")
+    
+    args = parser.parse_args()
+    
+    model_type = "thermal" if args.modality.lower() == "thermal" else "rgb"
+
+    if not args.weights_path:
+        if args.modality == "rgb":
+            args.weights_path = "./Firefly_RGB/model.safetensors"
+        elif args.modality == "thermal":
+            args.weights_path = "./Firefly_Thermal/model.safetensors"
+        print(f"No weights path provided, defaulting to: {args.weights_path}")
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    device = torch.device("cuda")
+
+    df = pd.read_csv(args.class_dict_path)
+    df = df.iloc[1:].reset_index(drop=True)
+    classes = df["name"]
+    id2label = classes.to_dict()
+    label2id = {v: k for k, v in id2label.items()}
+    
+    num_classes = len(classes)
+    print(f"Loaded {num_classes} classes from {args.class_dict_path}")
+    print(label2id)
+
+    config_dict = {
+        "num_classes": num_classes,
+        "ignore_label": 255,
+        "class_dict_path": args.class_dict_path,
+        "image_size": args.image_size,
+        "output_dir": args.output_dir
+    }
+
+    print("Loading dataset...")
+    if args.modality.lower() == "thermal":
+        val_set = InferenceDatasetThermal(
+            _root_dir=args.data_dir,
+            config=config_dict,
+        )
+    else:
+        val_set = InferenceDataset(
+            _root_dir=args.data_dir,
+            config=config_dict,
+        )
+
+    print(f"Number of images in validation set: {len(val_set)}")
+
+    weights_dir = os.path.dirname(args.weights_path) if os.path.isfile(args.weights_path) else args.weights_path
+    image_processor_loaded = False
+    
+    if weights_dir and os.path.exists(os.path.join(weights_dir, "preprocessor_config.json")):
+        try:
+            print(f"Loading image processor config from {weights_dir}...")
+            image_processor = AutoImageProcessor.from_pretrained(weights_dir)
+            image_processor_loaded = True
+        except Exception as e:
+            print(f"Warning: Failed to load image processor from {weights_dir}: {e}")
+
+    if not image_processor_loaded:
+        print("Falling back to standard image processor initialization...")
+        image_processor_args = {
+            "size": {"height": args.image_size, "width": args.image_size},
+            "crop_size": {"height": args.image_size, "width": args.image_size},
+            "image_mean": [0.485, 0.456, 0.406],
+            "image_std": [0.229, 0.224, 0.225],
+            "do_center_crop": False,
+            "do_normalize": True,
+            "do_resize": True,
+            "do_rescale": True,
+            "do_reduce_labels": True, 
+        }
+        try:
+            image_processor = AutoImageProcessor.from_pretrained("nvidia/mit-b3", **image_processor_args)
+        except:
+            image_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512", **image_processor_args)
+
+    print(f"Initializing {model_type} model...")
+    if model_type == "rgb":
+        dino_config = FireflyConfigRGB(
+            num_labels=num_classes,
+            image_size=args.image_size,
+            embedding_dim=256,
+            backbone_embed_dim=768,
+            patch_size=16,
+            repo_dir=args.dinov3_repo_dir,
+            model_name="dinov3_vitb16",
+            semantic_loss_ignore_index=255,
+        )
+        model = FireflyForSemanticSegmentationRGB(dino_config)
+
+    elif model_type == "thermal":
+        dino_config = FireflyConfigThermal(
+            num_labels=num_classes,
+            image_size=args.image_size,
+            embedding_dim=256,
+            backbone_embed_dim=768,
+            patch_size=16,
+            num_layers=12,
+            rein_token_length=100,
+            feature_layers=[2, 5, 8, 11],
+            repo_dir=args.dinov3_repo_dir,
+            model_name="dinov3_vitb16",
+            semantic_loss_ignore_index=255,
+        )
+        model = FireflyForSemanticSegmentationThermal(dino_config)
+
+    model.print_trainable_params()
+
+    if args.weights_path and os.path.exists(args.weights_path):
+        print(f"Loading weights from {args.weights_path}...")
+        try:
+            if args.weights_path.endswith('.safetensors'):
+                state_dict = load_file(args.weights_path)
+            else:
+                checkpoint = torch.load(args.weights_path, map_location='cpu')
+                if isinstance(checkpoint, dict):
+                    if "state_dict" in checkpoint:
+                        state_dict = checkpoint["state_dict"]
+                    elif "model" in checkpoint:
+                        state_dict = checkpoint["model"]
+                    elif "student_model" in checkpoint:
+                        state_dict = checkpoint["student_model"]
+                    else:
+                        state_dict = checkpoint
+                else:
+                    state_dict = checkpoint
+                    
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                if k.startswith("module."):
+                    new_state_dict[k[7:]] = v
+                else:
+                    new_state_dict[k] = v
+
+            msg = model.load_state_dict(new_state_dict, strict=False)
+            print(f"Weights Loaded. Missing keys: {len(msg.missing_keys)}, Unexpected keys: {len(msg.unexpected_keys)}")
+            if len(msg.missing_keys) > 0:
+                print(f"First 5 missing: {msg.missing_keys[:5]}")
+        except Exception as e:
+            print(f"Failed to load weights file: {e}")
+            raise
+    else:
+        print(f"Warning: No valid checkpoint found at {args.weights_path}. Proceeding with initialized weights.")
+
+    model.eval()
+    model.to(device)
+    
+    torch.backends.cudnn.benchmark = True
+    
+    loader = torch.utils.data.DataLoader(
+        val_set, 
+        batch_size=1,      
+        num_workers=4,    
+        pin_memory=True   
+    )
+
+    metric = evaluate.load("mean_iou")
+    
+    results, memory, global_gt_counts = generate_metrics_for_dataset(
+        loader, model, device, metric, config_dict, image_processor, visualize=args.visualize
+    )
+    
+    global_mean_iou = results["mean_iou"]
+    global_mean_acc = results["mean_accuracy"]
+    per_category_iou = results["per_category_iou"]
+
+    ground_truth_set = global_gt_counts
+    iou = np.array(per_category_iou)
+    iou = np.nan_to_num(iou, nan=0.0) 
+
+    gt_present_mask = (ground_truth_set > 0)
+    total_gt_freq = np.sum(ground_truth_set[gt_present_mask])
+
+    if total_gt_freq > 0:
+        fwIoU = np.sum(ground_truth_set[gt_present_mask] * iou[gt_present_mask]) / total_gt_freq
+    else:
+        fwIoU = 0.0
+
+    print("=" * 40)
+    print("EVALUATION RESULTS")
+    print("=" * 40)
+    print(f"Mean mIoU: {global_mean_iou:.4f}")
+    print(f"Freq Weighted IoU: {fwIoU:.4f}")
+    print(f"Mean pixel accuracy: {global_mean_acc:.4f}")
+    print("-" * 40)
+
+    for class_index, class_iou in enumerate(per_category_iou):
+        label = id2label.get(class_index, f"Class {class_index}")
+        print(f"{label}: Mean IoU = {class_iou:.4f}")
+
+    print("-" * 40)
+    print(f"Mean memory usage: {np.mean(memory):.2f} MB")
+    print(f"Max. GPU memory usage: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+
+    with open(os.path.join(args.output_dir, "per_class_scores.txt"), "w", encoding="utf-8") as f:
+        f.write(f"Mean mIoU: {global_mean_iou:.4f}\n")
+        f.write(f"Freq Weighted IoU: {fwIoU:.4f}\n")
+        f.write(f"Mean pixel accuracy: {global_mean_acc:.4f}\n")
+        f.write("-" * 40 + "\n")
+        for i, score in enumerate(per_category_iou):
+            label = id2label.get(i, f"Class {i}")
+            f.write(f"{label}: {score:.4f}\n")
+            
+    print(f"\nAll results saved to {args.output_dir}")

infer.py

@@ -0,0 +1,112 @@
+import os
+import argparse
+
+import cv2
+import numpy as np
+import torch
+from safetensors.torch import load_file
+from transformers import AutoImageProcessor
+
+from lib.utils_segfly import ID2COLOR, mask2label
+from lib.firefly_rgb import FireflyForSemanticSegmentationRGB, FireflyConfigRGB
+from lib.firefly_thermal import FireflyForSemanticSegmentationThermal, FireflyConfigThermal
+
+
+def run_infer(args):
+    if not args.weights_path:
+        if args.modality == "rgb":
+            args.weights_path = "./Firefly_RGB/model.safetensors"
+        else:
+            args.weights_path = "./Firefly_Thermal/model.safetensors"
+        print(f"No weights path provided, defaulting to: {args.weights_path}")
+
+    weights_dir = os.path.dirname(args.weights_path) if os.path.isfile(args.weights_path) else args.weights_path
+
+    print(f"Loading image processor from {weights_dir}...")
+    image_processor = AutoImageProcessor.from_pretrained(weights_dir)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    print(f"Initializing {args.modality} model...")
+    if args.modality == "rgb":
+        cfg = FireflyConfigRGB(
+            num_labels=15,
+            image_size=args.image_size,
+            embedding_dim=256,
+            backbone_embed_dim=768,
+            patch_size=16,
+            repo_dir=args.dinov3_repo_dir,
+            model_name="dinov3_vitb16",
+            semantic_loss_ignore_index=255,
+        )
+        model = FireflyForSemanticSegmentationRGB(cfg)
+    else:
+        cfg = FireflyConfigThermal(
+            num_labels=15,
+            image_size=args.image_size,
+            embedding_dim=256,
+            backbone_embed_dim=768,
+            patch_size=16,
+            num_layers=12,
+            rein_token_length=100,
+            feature_layers=[2, 5, 8, 11],
+            repo_dir=args.dinov3_repo_dir,
+            model_name="dinov3_vitb16",
+            semantic_loss_ignore_index=255,
+        )
+        model = FireflyForSemanticSegmentationThermal(cfg)
+
+    print(f"Loading weights from {args.weights_path}...")
+    state_dict = load_file(args.weights_path)
+    new_state_dict = {
+        k[7:] if k.startswith("module.") else k: v
+        for k, v in state_dict.items()
+    }
+    msg = model.load_state_dict(new_state_dict, strict=False)
+    print(f"Weights loaded. Missing: {len(msg.missing_keys)}, Unexpected: {len(msg.unexpected_keys)}")
+
+    model.eval()
+    model.to(device)
+
+    image_bgr = cv2.imread(args.image)
+    if image_bgr is None:
+        raise FileNotFoundError(f"Could not read image: {args.image}")
+    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    orig_h, orig_w = image_rgb.shape[:2]
+
+    inputs = image_processor(images=image_rgb, return_tensors="pt")
+    pixel_values = inputs.pixel_values.to(device, dtype=torch.float32)
+
+    with torch.no_grad():
+        logits = model(pixel_values).logits
+        upsampled = torch.nn.functional.interpolate(
+            logits.float(), size=(orig_h, orig_w), mode="bilinear", align_corners=False
+        )
+        pred = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()
+
+    os.makedirs(args.output, exist_ok=True)
+
+    pred_color = mask2label(pred, ID2COLOR)
+    pred_bgr = cv2.cvtColor(pred_color, cv2.COLOR_RGB2BGR)
+
+    stem = os.path.splitext(os.path.basename(args.image))[0]
+    out_path = os.path.join(args.output, f"{stem}_pred.png")
+    cv2.imwrite(out_path, pred_bgr)
+    print(f"Segmentation saved to {out_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Single-image inference for Firefly RGB / Thermal models")
+    parser.add_argument("--image", type=str, required=True, help="Path to the input image.")
+    parser.add_argument("--modality", type=str, default="rgb", choices=["rgb", "thermal"],
+                        help="Model modality: 'rgb' or 'thermal'.")
+    parser.add_argument("--output", type=str, default="./infer_output",
+                        help="Directory to save the colorized segmentation output.")
+    parser.add_argument("--weights_path", type=str, default="",
+                        help="Path to model weights (.safetensors). Auto-detected from modality if not set.")
+    parser.add_argument("--dinov3_repo_dir", type=str, default="./dinov3",
+                        help="Path to the local DINOv3 repository.")
+    parser.add_argument("--image_size", type=int, default=640,
+                        help="Resolution at which the image is fed to the model.")
+    args = parser.parse_args()
+    run_infer(args)

lib/firefly_rgb.py

@@ -0,0 +1,220 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SemanticSegmenterOutput
+
+# ==============================================================================
+# SHARED UTILITIES & LOSS FUNCTIONS
+# ==============================================================================
+
+def dice_loss(pred, target, skip_classes=None, eps=1e-6):
+    N, C = pred.shape
+    if N == 0:
+        return pred.sum() * 0.0
+
+    p = torch.softmax(pred, dim=1)
+    t = torch.nn.functional.one_hot(target, C).float()
+
+    if skip_classes is None:
+        skip_classes = []
+
+    present = (t.sum(dim=0) > 0)
+    for sc in skip_classes:
+        if 0 <= sc < C:
+            present[sc] = False
+
+    if present.sum() == 0:
+        return pred.new_tensor(0.)
+
+    p_sel = p[:, present]
+    t_sel = t[:, present]
+
+    inter = (p_sel * t_sel).sum(dim=0)
+    union = p_sel.sum(dim=0) + t_sel.sum(dim=0)
+    dice = (2 * inter + eps) / (union + eps)
+
+    return 1 - dice.mean()
+
+
+class DiceLoss(nn.Module):
+    def __init__(self, skip_classes=None, eps=1e-6):
+        super(DiceLoss, self).__init__()
+        self.skip_classes = skip_classes
+        self.eps = eps
+
+    def forward(self, pred, target):
+        return dice_loss(pred, target, skip_classes=self.skip_classes, eps=self.eps)
+
+
+def pad_to_multiple(x, multiple):
+    _, _, h, w = x.shape
+    pad_h = (multiple - h % multiple) % multiple
+    pad_w = (multiple - w % multiple) % multiple
+    padding = (0, pad_w, 0, pad_h)
+    x_padded = F.pad(x, padding, mode='reflect')
+    return x_padded, padding
+
+
+def crop_to_shape(x, target_h, target_w):
+    return x[:, :, :target_h, :target_w]
+
+# ==============================================================================
+# BASE CONFIGURATION & MODEL
+# ==============================================================================
+
+class FireflyBaseConfig(PretrainedConfig):
+    """Base Configuration for all Firefly models."""
+    def __init__(
+        self,
+        num_labels: int = 2,
+        image_size: int = 640,
+        embedding_dim: int = 256,
+        backbone_embed_dim: int = 768,
+        patch_size: int = 16,
+        repo_dir: str = None,
+        model_name: str = "dinov3_vitb16",
+        weights_path: str = None,
+        semantic_loss_ignore_index: int = 255,
+        dropout_ratio: float = 0.1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_labels = num_labels
+        self.image_size = image_size
+        self.embedding_dim = embedding_dim
+        self.backbone_embed_dim = backbone_embed_dim
+        self.patch_size = patch_size
+        self.repo_dir = repo_dir
+        self.model_name = model_name
+        self.weights_path = weights_path
+        self.semantic_loss_ignore_index = semantic_loss_ignore_index
+        self.dropout_ratio = dropout_ratio
+
+
+class FireflyBaseModel(PreTrainedModel):
+    """
+    Base Model for Firefly. Handles backbone initialization,
+    weight loading, and parameter freezing.
+    """
+    def __init__(self, config: FireflyBaseConfig):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.patch_size = config.patch_size
+        
+        if config.repo_dir and config.model_name:
+            print(f"Loading backbone: {config.model_name} from {config.repo_dir}")
+            self.backbone = torch.hub.load(
+                config.repo_dir,
+                config.model_name,
+                source="local",
+                weights=config.weights_path,
+                pretrained=False
+            )
+        else:
+            raise ValueError("repo_dir and model_name must be specified in the config.")
+        
+        if config.weights_path and os.path.exists(config.weights_path):
+            print(f"Loading fine-tuned backbone weights from: {config.weights_path}")
+            checkpoint = torch.load(config.weights_path, map_location='cpu')
+            
+            if isinstance(checkpoint, dict):
+                if "student_model" in checkpoint:
+                    state_dict = checkpoint["student_model"]
+                elif "model" in checkpoint:
+                    state_dict = checkpoint["model"]
+                else:
+                    state_dict = checkpoint
+            else:
+                state_dict = checkpoint
+
+            new_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+            msg = self.backbone.load_state_dict(new_state_dict, strict=False)
+            print(f"Weight loading result: {msg}")
+        else:
+            print(f"Warning: Weights path '{config.weights_path}' not found or not provided. Backbone initialized randomly/default.")
+        
+        self.backbone.eval()
+        self._setup_trainable_params()
+
+    def _setup_trainable_params(self):
+        """Freeze backbone parameters initially."""
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+            
+    def print_trainable_params(self):
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        print(f"Total Parameters: {total / 1e6:.2f}M")
+        print(f"Trainable Parameters: {trainable / 1e6:.2f}M")
+        print(f"Trainable Ratio: {100 * trainable / total:.2f}%")
+
+# ==============================================================================
+# SEGMENTATION HEAD & MODEL
+# ==============================================================================
+
+class MLPHead(nn.Module):
+    def __init__(self, in_channels=768, hidden_dim=256, num_classes=12, dropout_prob=0.1):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+            nn.Conv2d(hidden_dim, num_classes, kernel_size=1)
+        )
+    def forward(self, x):
+        return self.model(x)
+
+class FireflyConfigRGB(FireflyBaseConfig):
+    model_type = "firefly-rgb"
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+class FireflyForSemanticSegmentationRGB(FireflyBaseModel):
+    config_class = FireflyConfigRGB
+
+    def __init__(self, config: FireflyConfigRGB):
+        super().__init__(config)
+        self.head = MLPHead(
+            in_channels=config.backbone_embed_dim,
+            hidden_dim=config.embedding_dim,
+            num_classes=config.num_labels,
+            dropout_prob=config.dropout_ratio
+        )
+        self.loss_fn = DiceLoss()
+
+    def _extract_features(self, pixel_values: torch.Tensor, h_pad: int, w_pad: int):
+        features_raw = self.backbone.get_intermediate_layers(pixel_values, n=1)
+        feat = features_raw[0]
+        if isinstance(feat, tuple): feat = feat[0]
+
+        if feat.ndim == 3:
+            B, N, Dim = feat.shape
+            H_grid, W_grid = h_pad // self.config.patch_size, w_pad // self.config.patch_size
+            num_spatial_tokens = H_grid * W_grid
+            if N > num_spatial_tokens:
+                feat = feat[:, -num_spatial_tokens:, :]
+            feat = feat.permute(0, 2, 1).view(B, Dim, H_grid, W_grid)
+        return feat
+
+    def forward(self, pixel_values: torch.Tensor, labels: torch.Tensor = None, **kwargs):
+        original_h, original_w = pixel_values.shape[-2:]
+        x_padded, padding = pad_to_multiple(pixel_values, self.config.patch_size)
+        pad_h, pad_w = x_padded.shape[-2], x_padded.shape[-1]
+        
+        features = self._extract_features(x_padded, pad_h, pad_w)
+        out_upscaled = F.interpolate(self.head(features), size=(pad_h, pad_w), mode='bilinear', align_corners=False)
+        logits = crop_to_shape(out_upscaled, original_h, original_w)
+        
+        loss = None
+        if labels is not None:
+            valid_mask = labels != self.config.semantic_loss_ignore_index
+            logits_masked = logits.permute(0, 2, 3, 1)[valid_mask]
+            loss = self.loss_fn(logits_masked, labels[valid_mask])
+        
+        return SemanticSegmenterOutput(loss=loss, logits=logits)

lib/firefly_thermal.py

@@ -0,0 +1,428 @@
+import os
+import math
+from functools import reduce
+from operator import mul
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import SemanticSegmenterOutput
+
+# ===================================================================
+# 1. SHARED UTILITIES & LOSS FUNCTIONS
+# ===================================================================
+
+def dice_loss(pred, target, skip_classes=None, eps=1e-6):
+    N, C = pred.shape
+    
+    if N == 0:
+        return pred.sum() * 0.0
+
+    p = torch.softmax(pred, dim=1)                        
+
+    t = torch.nn.functional.one_hot(target, C).float()    
+
+    if skip_classes is None:
+        skip_classes = []
+
+    present = (t.sum(dim=0) > 0)
+
+    for sc in skip_classes:
+        if 0 <= sc < C:
+            present[sc] = False
+
+    if present.sum() == 0:
+        return pred.new_tensor(0.)
+
+    p_sel = p[:, present]
+    t_sel = t[:, present]
+
+    inter = (p_sel * t_sel).sum(dim=0)                     
+    union = p_sel.sum(dim=0) + t_sel.sum(dim=0)         
+    dice = (2 * inter + eps) / (union + eps)              
+
+    return 1 - dice.mean()
+
+class DiceLoss(nn.Module):
+    def __init__(self, skip_classes=None, eps=1e-6):
+        super(DiceLoss, self).__init__()
+        self.skip_classes = skip_classes
+        self.eps = eps
+
+    def forward(self, pred, target):
+        return dice_loss(pred, target, skip_classes=self.skip_classes, eps=self.eps)
+
+def pad_to_multiple(x, multiple):
+    """Pads the input tensor so its spatial dimensions are multiples of 'multiple'."""
+    _, _, h, w = x.shape
+    pad_h = (multiple - h % multiple) % multiple
+    pad_w = (multiple - w % multiple) % multiple
+    padding = (0, pad_w, 0, pad_h)  
+    x_padded = F.pad(x, padding, mode='reflect')
+    return x_padded, padding
+
+def crop_to_shape(x, target_h, target_w):
+    """Crops the input tensor back to the target spatial dimensions."""
+    return x[:, :, :target_h, :target_w]
+
+class MLPHead(nn.Module):
+    def __init__(self, in_channels=768, hidden_dim=256, num_classes=12, dropout_prob=0.1):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+            nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Dropout2d(p=dropout_prob),
+            nn.Conv2d(hidden_dim, num_classes, kernel_size=1)
+        )
+        self.upscale = True
+
+    def forward(self, x):
+        return self.model(x)
+
+# ===================================================================
+# 2. REIN IMPLEMENTATION
+# ===================================================================
+
+class Reins(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        embed_dims: int,
+        patch_size: int,
+        query_dims: int = 256,
+        token_length: int = 100,
+        use_softmax: bool = True,
+        scale_init: float = 0.001,
+    ) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed_dims = embed_dims
+        self.patch_size = patch_size
+        self.query_dims = query_dims
+        self.token_length = token_length
+        self.scale_init = scale_init
+        self.use_softmax = use_softmax
+        self.create_model()
+
+    def create_model(self):
+        self.learnable_tokens = nn.Parameter(
+            torch.empty([self.num_layers, self.token_length, self.embed_dims])
+        )
+        self.scale = nn.Parameter(torch.tensor(self.scale_init))
+        
+        self.mlp_token2feat = nn.Linear(self.embed_dims, self.embed_dims)
+        self.mlp_delta_f = nn.Linear(self.embed_dims, self.embed_dims)
+        
+        val = math.sqrt(
+            6.0
+            / float(
+                3 * reduce(mul, (self.patch_size, self.patch_size), 1) + self.embed_dims
+            )
+        )
+        nn.init.uniform_(self.learnable_tokens.data, -val, val)
+        nn.init.kaiming_uniform_(self.mlp_delta_f.weight, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.mlp_token2feat.weight, a=math.sqrt(5))
+
+        self.transform = nn.Linear(self.embed_dims, self.query_dims)
+        self.merge = nn.Linear(self.query_dims * 3, self.query_dims)
+
+    def get_tokens(self, layer: int) -> Tensor:
+        if layer == -1:
+            return self.learnable_tokens
+        else:
+            return self.learnable_tokens[layer]
+
+    def forward(
+        self, feats: Tensor, layer: int, batch_first=False, has_cls_token=True
+    ) -> Tensor:
+        if batch_first:
+            feats = feats.permute(1, 0, 2) 
+        
+        if has_cls_token:
+            cls_token, feats = torch.tensor_split(feats, [1], dim=0)
+            
+        tokens = self.get_tokens(layer)
+        
+        delta_feat = self.forward_delta_feat(
+            feats,
+            tokens,
+            layer,
+        )
+        
+        delta_feat = delta_feat * self.scale
+        feats = feats + delta_feat 
+        
+        if has_cls_token:
+            feats = torch.cat([cls_token, feats], dim=0)
+            
+        if batch_first:
+            feats = feats.permute(1, 0, 2)
+            
+        return feats
+
+    def forward_delta_feat(self, feats: Tensor, tokens: Tensor, layers: int) -> Tensor:
+        attn = torch.einsum("nbc,mc->nbm", feats, tokens)
+        
+        if self.use_softmax:
+            attn = attn * (self.embed_dims**-0.5)
+            attn = F.softmax(attn, dim=-1)
+            
+        delta_f = torch.einsum(
+            "nbm,mc->nbc",
+            attn[:, :, 1:], 
+            self.mlp_token2feat(tokens[1:, :]),
+        )
+        
+        delta_f = self.mlp_delta_f(delta_f + feats)
+        return delta_f
+
+class ReinHook:
+    def __init__(self, rein_module, layer_idx):
+        self.rein_module = rein_module
+        self.layer_idx = layer_idx
+
+    def __call__(self, module, input, output):
+        is_tuple = isinstance(output, tuple)
+        is_list = isinstance(output, list)
+        
+        if is_tuple:
+            x = output[0]
+        elif is_list:
+            x = output[0]
+        else:
+            x = output
+
+        refined_x = self.rein_module(x, self.layer_idx, batch_first=True, has_cls_token=True)
+        
+        if is_tuple:
+            return (refined_x,) + output[1:]
+        elif is_list:
+            output[0] = refined_x
+            return output
+        else:
+            return refined_x
+
+class FireflyConfigThermal(PretrainedConfig):
+    model_type = "firefly-thermal"
+
+    def __init__(
+        self,
+        num_labels: int = 2,
+        image_size: int = 512,
+        embedding_dim: int = 256,
+        backbone_embed_dim: int = 768,
+        patch_size: int = 16,
+        num_layers: int = 12,
+        rein_token_length: int = 100,  
+        feature_layers: list = None,
+        repo_dir: str = None,
+        model_name: str = "dinov3_vitb16",
+        backbone_weights_path: str = None, 
+        finetuned_weights_path: str = None, 
+        semantic_loss_ignore_index: int = 255,
+        dropout_ratio: float = 0.1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_labels = num_labels
+        self.image_size = image_size
+        self.embedding_dim = embedding_dim
+        self.backbone_embed_dim = backbone_embed_dim
+        self.patch_size = patch_size
+        self.num_layers = num_layers
+        self.rein_token_length = rein_token_length
+        self.feature_layers = feature_layers or [2, 5, 8, 11]
+        self.repo_dir = repo_dir
+        self.model_name = model_name
+        self.backbone_weights_path = backbone_weights_path
+        self.finetuned_weights_path = finetuned_weights_path
+        self.semantic_loss_ignore_index = semantic_loss_ignore_index
+        self.dropout_ratio = dropout_ratio
+
+class FireflyForSemanticSegmentationThermal(PreTrainedModel):
+    config_class = FireflyConfigThermal
+
+    def __init__(self, config: FireflyConfigThermal):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+        
+        print(f"Loading backbone architecture: {config.model_name}")
+        self.backbone = torch.hub.load(
+            config.repo_dir,
+            config.model_name,
+            source="local",
+            weights=config.backbone_weights_path,
+            pretrained=False
+        )
+    
+        print(f"Injecting Rein Module (Tokens={config.rein_token_length})...")
+        self.rein = Reins(
+            num_layers=config.num_layers,
+            embed_dims=config.backbone_embed_dim,
+            patch_size=config.patch_size,
+            token_length=config.rein_token_length
+        )
+
+        if hasattr(self.backbone, 'blocks'):
+            blocks = self.backbone.blocks
+        elif hasattr(self.backbone, 'transformer') and hasattr(self.backbone.transformer, 'blocks'): 
+             blocks = self.backbone.transformer.blocks
+        else:
+            raise AttributeError("Could not find '.blocks' in backbone model. Check model structure.")
+
+        self.hooks = []
+        for i, block in enumerate(blocks):
+            hook_fn = ReinHook(self.rein, i)
+            handle = block.register_forward_hook(hook_fn)
+            self.hooks.append(handle)
+        
+        if config.finetuned_weights_path and os.path.exists(config.finetuned_weights_path):
+            print(f"Loading finetuned rein adapter weights: {config.finetuned_weights_path}")
+            checkpoint = torch.load(config.finetuned_weights_path, map_location='cpu')
+            
+            state_dict = checkpoint
+            if isinstance(checkpoint, dict):
+                if "rein_model" in checkpoint:
+                    state_dict = checkpoint["rein_model"]
+                elif "student_model" in checkpoint:
+                    state_dict = checkpoint["student_model"]
+                elif "state_dict" in checkpoint:
+                    state_dict = checkpoint["state_dict"]
+                elif "model" in checkpoint:
+                    state_dict = checkpoint["model"]
+
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                clean_k = k
+                for prefix in ["module.", "rein_model.", "rein.", "student_model.", "model."]:
+                    if clean_k.startswith(prefix):
+                        clean_k = clean_k.replace(prefix, "", 1)
+                new_state_dict[clean_k] = v
+
+            msg = self.rein.load_state_dict(new_state_dict, strict=False)
+            print(msg)
+            if len(msg.missing_keys) > 0:
+                print(f"[Diagnostic] Found {len(msg.missing_keys)} missing keys in Rein module loading.")
+                print(f"[Warning] Missing keys: {msg.missing_keys}")
+            else:
+                print("[Success] All Rein Adapter weights perfectly matched and loaded!")
+        else:
+            print(f"Warning: Finetuned weights not found at {config.finetuned_weights_path}. Using base/random weights.")
+
+        self.head = MLPHead(
+            in_channels=config.backbone_embed_dim,
+            hidden_dim=config.embedding_dim,
+            num_classes=config.num_labels,
+            dropout_prob=config.dropout_ratio
+        )
+        
+        self.loss_fn = DiceLoss()
+        
+        self.feature_layers = config.feature_layers
+        self._setup_trainable_params()
+
+    def _setup_trainable_params(self):
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        
+        for param in self.rein.parameters():
+            param.requires_grad = True
+        
+        for param in self.head.parameters():
+            param.requires_grad = True
+
+    def _extract_features(self, pixel_values: torch.Tensor, h_pad: int, w_pad: int):
+        features_raw = self.backbone.get_intermediate_layers(
+            pixel_values, 
+            n=1 
+        )
+        
+        feat = features_raw[0]
+        if isinstance(feat, tuple):
+            feat = feat[0]
+
+        if feat.ndim == 3:
+            B, N, Dim = feat.shape
+            H_grid = h_pad // self.config.patch_size
+            W_grid = w_pad // self.config.patch_size
+            
+            num_spatial_tokens = H_grid * W_grid
+            if N > num_spatial_tokens:
+                feat = feat[:, -num_spatial_tokens:, :]
+                
+            feat = feat.permute(0, 2, 1).view(B, Dim, H_grid, W_grid)
+            
+        return feat
+
+    def forward(self, pixel_values: torch.Tensor, labels: torch.Tensor = None, **kwargs):
+        original_h, original_w = pixel_values.shape[-2:]
+        
+        x_padded, padding = pad_to_multiple(pixel_values, self.config.patch_size)
+        pad_h, pad_w = x_padded.shape[-2], x_padded.shape[-1]
+        
+        features = self._extract_features(x_padded, pad_h, pad_w)
+        out = self.head(features)
+        
+        out_upscaled = F.interpolate(
+            out, 
+            size=(pad_h, pad_w), 
+            mode='bilinear', 
+            align_corners=False
+        )
+        
+        logits = crop_to_shape(out_upscaled, original_h, original_w)
+        
+        loss = None
+        if labels is not None:
+            valid_mask = labels != self.config.semantic_loss_ignore_index
+            logits_masked = logits.permute(0, 2, 3, 1)[valid_mask]
+            labels_masked = labels[valid_mask]
+            loss = self.loss_fn(logits_masked, labels_masked)
+        
+        return SemanticSegmenterOutput(
+            loss=loss,
+            logits=logits,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path=None, **kwargs):
+        config_kwargs = {
+            "num_labels": kwargs.pop("num_labels", 2),
+            "image_size": kwargs.pop("image_size", 512),
+            "embedding_dim": kwargs.pop("embedding_dim", 256),
+            "backbone_embed_dim": kwargs.pop("backbone_embed_dim", 768),
+            "patch_size": kwargs.pop("patch_size", 16),
+            "num_layers": kwargs.pop("num_layers", 12),
+            "rein_token_length": kwargs.pop("rein_token_length", 100),
+            "feature_layers": kwargs.pop("feature_layers", [2, 5, 8, 11]),
+            "repo_dir": kwargs.pop("repo_dir", None),
+            "model_name": kwargs.pop("model_name", "dinov3_vitb16"),
+            "backbone_weights_path": kwargs.pop("backbone_weights_path", ""),
+            "finetuned_weights_path": kwargs.pop("finetuned_weights_path", pretrained_model_name_or_path),
+            "semantic_loss_ignore_index": kwargs.pop("semantic_loss_ignore_index", 255),
+            "dropout_ratio": kwargs.pop("dropout_ratio", 0.1),
+            "id2label": kwargs.pop("id2label", None),
+            "label2id": kwargs.pop("label2id", None),
+        }
+        
+        kwargs.pop("ignore_mismatched_sizes", None)
+        config = kwargs.pop("config", None)
+    
+        if config is None:
+            config = FireflyConfigThermal(**config_kwargs)
+            
+        model = cls(config)
+        return model
+
+    def print_trainable_params(self):
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        print(f"Total Parameters: {total / 1e6:.2f}M")
+        print(f"Trainable Parameters: {trainable / 1e6:.2f}M")
+        print(f"Trainable Ratio: {100 * trainable / total:.2f}%")

lib/utils_segfly.py

@@ -0,0 +1,251 @@
+"""Utility functions for dataset loading, visualization, and timing."""
+
+import os
+from functools import wraps
+from time import perf_counter
+from typing import List
+
+import cv2
+import numpy as np
+import pandas as pd
+import torch
+import albumentations as aug
+from torch.utils.data import Dataset
+
+
+# ==============================================================================
+# VISUALIZATION PALETTE
+# ==============================================================================
+
+ID2COLOR = {
+    0:   [128, 0, 128],
+    1:   [204, 163, 72],
+    2:   [128, 0, 0],
+    3:   [192, 192, 192],
+    4:   [0, 255, 0],
+    5:   [112, 148, 32],
+    6:   [64, 64, 0],
+    7:   [255, 255, 0],
+    8:   [0, 128, 128],
+    9:   [0, 0, 255],
+    10:  [255, 0, 0],
+    11:  [64, 160, 120],
+    12:  [128, 64, 128],
+    13:  [240, 120, 120],
+    14:  [128, 128, 64],
+    255: [0, 0, 0],
+}
+
+
+def mask2label(mask, palette):
+    """Convert a class-index mask to a colorized RGB image."""
+    if mask.ndim == 3 and mask.shape[0] == 1:
+        mask = mask.squeeze(0)
+    if isinstance(mask, torch.Tensor):
+        mask = mask.detach().cpu().numpy()
+    
+    color_seg = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
+    for label, color in palette.items():
+        color_seg[mask == label, :] = color
+    return color_seg
+
+
+# ==============================================================================
+# DATASET CLASSES
+# ==============================================================================
+
+class InferenceDataset(Dataset):
+    """Dataset for RGB inference with class remapping and dropping."""
+
+    def __init__(
+        self,
+        _root_dir: str,
+        config: dict,
+        _ignore_index: int = 255,
+        _original_num_classes: int = 22,
+    ):
+        """
+        Args:
+            _root_dir (str): Path to the root directory. Walked recursively;
+                images are collected from any subfolder named `src/` and masks
+                from any subfolder named `gt/`. No specific top-level structure
+                is required.
+            config (dict): Dictionary containing the configuration.
+            _ignore_index (int): Index used for dropped classes.
+            _original_num_classes (int): Number of original classes.
+        """
+        self.root_dir = _root_dir
+        self.transforms = aug.Compose(
+            [aug.Resize(config["image_size"], config["image_size"], interpolation=cv2.INTER_NEAREST)]
+        )
+        self.ignore_index = _ignore_index  
+        self.mapping_lut = np.full(256, self.ignore_index, dtype=np.uint8)
+
+        df = pd.read_csv(config["class_dict_path"])
+
+        new_class_id = 0
+        raw_to_new_dict = {}
+        
+        for _, row in df.iterrows():
+            raw_id = int(row['id'])
+            if raw_id == self.ignore_index:
+                continue
+            self.mapping_lut[raw_id] = new_class_id
+            raw_to_new_dict[raw_id] = new_class_id
+            new_class_id += 1
+
+        if 9 in raw_to_new_dict:
+            target_id_for_merges = raw_to_new_dict[9]
+            self.mapping_lut[5] = target_id_for_merges
+            self.mapping_lut[22] = target_id_for_merges
+        
+        self.num_new_classes = new_class_id
+
+        image_file_names = []
+        annotation_file_names = []
+        for root, _, files in os.walk(self.root_dir):
+            for name in files:
+                if ".ipynb_checkpoints" in root:
+                    continue
+                if "src" in root and name.lower().endswith((".png", ".jpeg", ".jpg")):
+                    image_file_names.append(os.path.join(root, name))
+                elif "gt" in root and name.lower().endswith(".png"):
+                    if "vis_rem" not in os.path.join(root, name):
+                        annotation_file_names.append(os.path.join(root, name))
+
+        self.images = sorted(image_file_names)
+        self.annotations = sorted(annotation_file_names)
+
+        assert len(self.images) == len(self.annotations), (
+            f"There must be as many images as there are segmentation maps. "
+            f"Image = {len(self.images)}, Masks = {len(self.annotations)}"
+        )
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image = cv2.imread(self.images[idx])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        segmentation_map = cv2.imread(self.annotations[idx], cv2.IMREAD_GRAYSCALE)
+
+        segmentation_map = self.mapping_lut[segmentation_map]
+
+        if self.transforms is not None:
+            augmented = self.transforms(image=image, mask=segmentation_map)
+            return augmented["image"], augmented["mask"]
+
+        return image, segmentation_map
+
+
+class InferenceDatasetThermal(Dataset):
+    """Dataset for thermal inference with class remapping and dropping."""
+
+    def __init__(
+        self,
+        _root_dir: str,
+        config: dict,
+        _ignore_index: int = 255,
+    ):
+        """
+        Args:
+            _root_dir (str): Path to the root directory. Walked recursively;
+                images are collected from any subfolder named `thermal_src/`
+                and masks from any subfolder named `gt/`. No specific top-level
+                structure is required.
+            config (dict): Dictionary containing the configuration.
+            _ignore_index (int): Index used for dropped classes.
+        """
+        self.root_dir = _root_dir
+        self.transforms = aug.Compose(
+            [aug.Resize(config["image_size"], config["image_size"], interpolation=cv2.INTER_NEAREST)]
+        )
+        self.ignore_index = _ignore_index  
+        self.mapping_lut = np.full(256, self.ignore_index, dtype=np.uint8)
+
+        df = pd.read_csv(config["class_dict_path"])
+
+        new_class_id = 0
+        raw_to_new_dict = {}
+        
+        for _, row in df.iterrows():
+            raw_id = int(row['id'])
+            self.mapping_lut[raw_id] = new_class_id
+            raw_to_new_dict[raw_id] = new_class_id
+            new_class_id += 1
+
+        if 9 in raw_to_new_dict:
+            target_id_for_merges = raw_to_new_dict[9]
+            self.mapping_lut[5] = target_id_for_merges
+            self.mapping_lut[22] = target_id_for_merges
+        
+        self.num_new_classes = new_class_id
+
+        image_file_names = []
+        annotation_file_names = []
+        for root, _, files in os.walk(self.root_dir):
+            for name in files:
+                if ".ipynb_checkpoints" in root:
+                    continue
+                if "thermal_src" in root and name.lower().endswith((".png", ".jpeg", ".jpg")):
+                    image_file_names.append(os.path.join(root, name))
+                elif "gt" in root and name.lower().endswith(".png"):
+                    annotation_file_names.append(os.path.join(root, name))
+
+        self.images = sorted(image_file_names)
+        self.annotations = sorted(annotation_file_names)
+
+        assert len(self.images) == len(self.annotations), (
+            f"There must be as many images as there are segmentation maps."
+            f"Image = {len(self.images)}, Masks = {len(self.annotations)}"
+        )
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image = cv2.imread(self.images[idx])
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        segmentation_map = cv2.imread(self.annotations[idx], cv2.IMREAD_GRAYSCALE)
+
+        segmentation_map = self.mapping_lut[segmentation_map]
+
+        if self.transforms is not None:
+            augmented = self.transforms(image=image, mask=segmentation_map)
+            return augmented["image"], augmented["mask"]
+
+        return image, segmentation_map
+
+
+# ==============================================================================
+# TIMING UTILITY
+# ==============================================================================
+
+class Timing:
+    """Class to time functions and methods."""
+
+    def __init__(self):
+        self.inf_time = []
+
+    def __call__(self, func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start = perf_counter()
+            result = func(*args, **kwargs)
+            end = perf_counter()
+            self.inf_time.append(end - start)
+            return result
+
+        return wrapper
+
+    def print_average(self):
+        """Print the average timing of the function."""
+        print(f"---\tFunction took {np.mean(self.inf_time):.4f} seconds to run!")
+
+    def get_average(self) -> float:
+        """Return the average timing of the function."""
+        return float(np.mean(self.inf_time))
+
+    def reset(self):
+        """Reset the timing of the function."""
+        self.inf_time = []

requirements.txt

@@ -0,0 +1,15 @@
+torch>=2.9
+transformers==4.37.2
+evaluate>=0.4.0
+datasets>=2.20.0
+numpy>=1.24.4
+pandas>=1.5.3
+albumentations>=1.3.1
+opencv-python-headless>=4.9.0
+safetensors
+tqdm
+pyyaml
+matplotlib>=3.8.2
+torchmetrics>=1.8.2
+torchvision
+termcolor