initial commit

app.py

@@ -0,0 +1,105 @@
+import streamlit as st
+import numpy as np
+import os
+import pathlib
+from inference import infer, InferenceModel
+
+# -----------------------------------------------------------------------------
+# class SatvisionDemoApp
+#
+# Directory Structure: base-directory/MOD09GA/year
+#                                     MOD09GQ/year
+#                                     MYD09GA/year
+#                                     MYD09GQ/year
+#
+# -----------------------------------------------------------------------------
+class SatvisionDemoApp:
+
+    # -------------------------------------------------------------------------
+    # __init__
+    # -------------------------------------------------------------------------
+    def __init__(self):
+
+        self.thumbnail_dir = pathlib.Path('data/thumbnails')
+        self.image_dir = pathlib.Path('data/images')
+        print(self.thumbnail_dir)
+        self.thumbnail_files = sorted(list(self.thumbnail_dir.glob('sv-*.png')))
+        self.image_files = sorted(list(self.image_dir.glob('sv-*.npy')))
+        print(list(self.image_files))
+        self.thumbnail_names = [str(tn_path.name) for tn_path in self.thumbnail_files]
+        print(self.thumbnail_names)
+
+        self.inferenceModel =  InferenceModel()
+
+    # -------------------------------------------------------------------------
+    # render_sidebar
+    # -------------------------------------------------------------------------
+    def render_sidebar(self):
+
+        st.sidebar.header("Select an Image")
+
+        for index, thumbnail in enumerate(self.thumbnail_names):
+
+            thumbnail_path = self.thumbnail_dir / thumbnail
+
+            # thumbnail_arr = np.load(thumbnail_path)
+            print(str(thumbnail_path))
+
+            st.sidebar.image(str(thumbnail_path), use_column_width=True, caption=thumbnail)
+
+    # -------------------------------------------------------------------------
+    # render_main_app
+    # -------------------------------------------------------------------------
+    def render_main_app(self):
+
+        st.title("Satvision-Base Demo")
+
+        st.header("Image Reconstruction Process")
+        selected_image_index = st.sidebar.selectbox(
+                    "Select an Image",
+                    self.thumbnail_names)
+        print(selected_image_index)
+
+        selected_image = self.load_selected_image(selected_image_index)
+
+        image, masked_input, output = self.inferenceModel.infer(selected_image)
+
+        col1, col2, col3 = st.columns(3, gap="large")
+
+        # Display the selected image with a title three times side-by-side
+
+        with col1:
+            st.image(image, use_column_width=True, caption="Input")
+
+        with col2:
+            st.image(masked_input, use_column_width=True, caption="Input Masked")
+
+        with col3:
+            st.image(output, use_column_width=True, caption="Reconstruction")
+
+    # -------------------------------------------------------------------------
+    # load_selected_image
+    # -------------------------------------------------------------------------
+    def load_selected_image(self, image_name):
+
+        # Load the selected image using NumPy (replace this with your image loading code)
+        image_name = image_name.replace('.png', '.npy')
+
+        image = np.load(self.image_dir / image_name)
+        image = np.moveaxis(image, 0, 2)
+        return image
+
+# -----------------------------------------------------------------------------
+# main
+# -----------------------------------------------------------------------------
+def main():
+
+    app = SatvisionDemoApp()
+
+    app.render_main_app()
+
+    app.render_sidebar()
+
+if __name__ == "__main__":
+
+    main()

ckpt_epoch_800.pth

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

data/images/sv-demo-mod09ga-11.npy

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

data/images/sv-demo-mod09ga-12.npy

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

data/images/sv-demo-mod09ga-6.npy

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

data/images/sv-demo-mod09ga-7.npy

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

inference.py

@@ -0,0 +1,227 @@
+import numpy as np
+
+import torch
+import joblib
+import numpy as np
+
+import torchvision.transforms as T
+import sys
+
+sys.path.append('pytorch-caney')
+# from pytorch_caney.models.mim.mim import build_mim_model
+
+
+class Transform:
+    """
+    torchvision transform which transforms the input imagery into
+    addition to generating a MiM mask
+    """
+
+    def __init__(self, config):
+
+        self.transform_img = \
+            T.Compose([
+                T.ToTensor(),
+                T.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE)),
+            ])
+
+        model_patch_size = config.MODEL.SWINV2.PATCH_SIZE
+
+        self.mask_generator = SimmimMaskGenerator(
+            input_size=config.DATA.IMG_SIZE,
+            mask_patch_size=config.DATA.MASK_PATCH_SIZE,
+            model_patch_size=model_patch_size,
+            mask_ratio=config.DATA.MASK_RATIO,
+        )
+
+    def __call__(self, img):
+
+        img = self.transform_img(img)
+        mask = self.mask_generator()
+
+        return img, mask
+
+
+class SimmimMaskGenerator:
+    """
+    Generates the masks for masked-image-modeling
+    """
+    def __init__(self,
+                 input_size=192,
+                 mask_patch_size=32,
+                 model_patch_size=4,
+                 mask_ratio=0.6):
+        self.input_size = input_size
+        self.mask_patch_size = mask_patch_size
+        self.model_patch_size = model_patch_size
+        self.mask_ratio = mask_ratio
+
+        assert self.input_size % self.mask_patch_size == 0
+        assert self.mask_patch_size % self.model_patch_size == 0
+
+        self.rand_size = self.input_size // self.mask_patch_size
+        self.scale = self.mask_patch_size // self.model_patch_size
+
+        self.token_count = self.rand_size ** 2
+        self.mask_count = int(np.ceil(self.token_count * self.mask_ratio))
+
+    def __call__(self):
+        mask = self.make_simmim_mask(self.token_count, self.mask_count,
+                                self.rand_size, self.scale)
+        mask = mask.repeat(self.scale, axis=0).repeat(self.scale, axis=1)
+        return mask
+    
+    @staticmethod
+    def make_simmim_mask(token_count, mask_count, rand_size, scale):
+        """JIT-compiled random mask generation
+
+        Args:
+            token_count
+            mask_count
+            rand_size
+            scale
+
+        Returns:
+            mask
+        """
+        mask_idx = np.random.permutation(token_count)[:mask_count]
+        mask = np.zeros(token_count, dtype=np.int64)
+        mask[mask_idx] = 1
+        mask = mask.reshape((rand_size, rand_size))
+        return mask
+
+
+class InferenceModel(object):
+
+    def __init__(self):
+        self.checkpoint_path = 'ckpt_epoch_800.pth'
+        self.config_path = 'simmim_pretrain__satnet_swinv2_base__img192_window12__800ep_v3_no_norm.config.sav'
+        self.architecture_path = 'model.sav'
+
+        self.config = joblib.load(self.config_path)
+        self.model = joblib.load(self.architecture_path)
+        self.load_checkpoint()
+
+        self.transform = Transform(self.config)
+
+
+    def load_checkpoint(self):
+
+
+        checkpoint = torch.load(self.checkpoint_path, map_location='cpu')
+
+        # re-map keys due to name change (only for loading provided models)
+        rpe_mlp_keys = [k for k in checkpoint['model'].keys() if "rpe_mlp" in k]
+
+        for k in rpe_mlp_keys:
+
+            checkpoint['model'][k.replace(
+                'rpe_mlp', 'cpb_mlp')] = checkpoint['model'].pop(k)
+
+        msg = self.model.load_state_dict(checkpoint['model'], strict=False)
+
+        print(msg)
+
+        del checkpoint
+
+        torch.cuda.empty_cache()
+
+    @staticmethod
+    def minmax_norm(img_arr):
+        arr_min = img_arr.min()
+        arr_max = img_arr.max()
+        img_arr_scaled = (img_arr - arr_min) / (arr_max - arr_min)
+        img_arr_scaled = img_arr_scaled * 255
+        img_arr_scaled = img_arr_scaled.astype(np.uint8)
+        return img_arr_scaled
+
+    # -------------------------------------------------------------------------
+    # load_selected_image
+    # -------------------------------------------------------------------------
+    def preprocess(self, image):
+
+        image, mask = self.transform(image)
+
+        image = image.unsqueeze(0)
+    
+        mask = torch.tensor(mask).unsqueeze(0)
+
+        print(image.size())
+        print(mask.shape)
+
+        return image, mask
+
+    # -------------------------------------------------------------------------
+    # load_selected_image
+    # -------------------------------------------------------------------------
+    def predict(self, image, mask):
+
+        with torch.no_grad():
+    
+            logits = self.model.encoder(image, mask)
+    
+            image_recon = self.model.decoder(logits)
+        
+        image_recon = image_recon.numpy()[0, :, :, :]
+
+        return image_recon
+
+    # -------------------------------------------------------------------------
+    # load_selected_image
+    # -------------------------------------------------------------------------
+    @staticmethod
+    def process_mask(mask):
+        mask = mask.repeat_interleave(4, 1).repeat_interleave(4, 2).unsqueeze(1).contiguous()
+        mask = mask[0, 0, :, :]
+        mask = np.stack([mask, mask, mask], axis=-1)
+        return mask
+
+    # -------------------------------------------------------------------------
+    # load_selected_image
+    # -------------------------------------------------------------------------
+    def infer(self, image):
+
+        image, mask = self.preprocess(image)
+
+        img_recon = self.predict(image, mask)
+
+        mask = self.process_mask(mask)
+
+        img_normed = self.minmax_norm(image.numpy()[0, :, :, :])
+
+        print(img_normed.shape)
+        rgb_image = np.stack((img_normed[0, :, :],
+                            img_normed[3, :, :],
+                            img_normed[2, :, :]),
+                            axis=-1)
+
+        img_recon = self.minmax_norm(img_recon)
+        rgb_image_recon = np.stack((img_recon[0, :, :],
+                                    img_recon[3, :, :],
+                                    img_recon[2, :, :]),
+                                    axis=-1)
+
+        rgb_masked = np.where(mask == 0, rgb_image, rgb_image_recon)
+        rgb_image_masked = np.where(mask == 1, 0, rgb_image)
+        rgb_recon_masked = rgb_masked# self.minmax_norm(rgb_masked)
+
+        return rgb_image, rgb_image_masked, rgb_recon_masked
+
+
+def infer(array_input: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+
+    masked_input = np.random.rand(256, 256, 3)
+
+    output = np.random.rand(256, 256, 3)
+
+    return masked_input, output
+
+if __name__ == '__main__':
+    inferenceModel = InferenceModel()
+
+    image = np.load('data/images/sv-demo-mod09ga-11.npy')
+    print(image.shape)
+    image = np.moveaxis(image, 0, 2)
+    print(image.shape)
+
+    inference = inferenceModel.infer(image)

model.sav

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

pytorch-caney/.readthedocs.yaml

@@ -0,0 +1,9 @@
+version: 2
+
+build:
+  os: "ubuntu-20.04"
+  tools:
+    python: "3.8"
+
+sphinx:
+  fail_on_warning: true

pytorch-caney/CODE_OF_CONDUCT.md

@@ -0,0 +1,128 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+support@nccs.nasa.gov.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

pytorch-caney/CONTRIBUTING.md

@@ -0,0 +1,173 @@
+# Contributing
+
+When contributing to this repository, please first discuss the change you wish to make via issue,
+email, or any other method with the owners of this repository before making a change.
+
+Please note we have a code of conduct, please follow it in all your interactions with the project.
+
+## Pull Request Process
+
+1. Ensure any install or build dependencies are removed before the end of the layer when doing a
+   build.
+2. Update the README.md with details of changes to the interface, this includes new environment
+   variables, exposed ports, useful file locations and container parameters.
+3. Increase the version numbers in any examples files and the README.md to the new version that this
+   Pull Request would represent. The versioning scheme we use is [SemVer](http://semver.org/).
+4. Regenerate any additional documentation using PDOC (usage details listed below).
+5. Document the proposed changes in the CHANGELOG.md file.
+6. You may submit your merge request for review and the change will be reviewed.
+
+## Code of Conduct
+
+### Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, gender identity and expression, level of experience,
+nationality, personal appearance, race, religion, or sexual identity and
+orientation.
+
+### Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+### Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+### Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+### Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at [INSERT EMAIL ADDRESS]. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+### Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at [http://contributor-covenant.org/version/1/4][version]
+
+[homepage]: http://contributor-covenant.org
+[version]: http://contributor-covenant.org/version/1/4/
+
+## Appendix
+
+### Generating Documentation
+
+This repository follows semi-automatic documentation generation. The following
+is an example of how to generate documentation for a single module.
+
+```bash
+conda activate tensorflow-caney
+pdoc --html tensorflow-caney/raster.py --force
+```
+
+### Linting
+
+This project uses flake8 for PREP8 linting and format. Every submodule should include
+a test section in the tests directory. Refer to the text directory for more examples.
+The Python unittests library is used for these purposes.
+
+### Documenting Methods
+
+The following documentation format should be followed below each method to allow for
+explicit semi-automatic documentation generation.
+
+```bash
+   """
+   Read raster and append data to existing Raster object
+   Args:
+      filename (str): raster filename to read from
+      bands (str list): list of bands to append to object, e.g ['Red']
+      chunks_band (int): integer to map object to memory, z
+      chunks_x (int): integer to map object to memory, x
+      chunks_y (int): integer to map object to memory, y
+   Return:
+      raster (raster object): raster object to manipulate rasters
+   ----------
+   Example
+   ----------
+      raster.readraster(filename, bands)
+   """
+```
+
+### Format of CHANGELOG
+
+The following describes the format for each CHANGELOG release. If there are no contributions
+in any of the sections, they are removed from the description.
+
+```bash
+## [0.0.3] - 2020-12-14
+
+### Added
+- Short description
+
+### Fixed
+- Short description
+
+### Changed
+- Short description
+
+### Removed
+- Short description
+
+### Approved
+Approver Name, Email
+```
+
+### Example Using Container in ADAPT
+
+```bash
+module load singularity
+singularity shell -B $your_mounts --nv tensorflow-caney
+```
+
+### Current Workflow
+
+```bash
+module load singularity
+singularity shell --nv -B /lscratch,/css,/explore/nobackup/projects/ilab,/explore/nobackup/people /explore/nobackup/projects/ilab/containers/tensorflow-caney-2022.12
+export PYTHONPATH="/explore/nobackup/people/jacaraba/development/tensorflow-caney:/adapt/nobackup/people/jacaraba/development/vhr-cnn-chm"
+```

pytorch-caney/LICENSE

@@ -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.

pytorch-caney/LICENSE.md

@@ -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.

pytorch-caney/README.md

@@ -0,0 +1,132 @@
+# pytorch-caney
+
+Python package for lots of Pytorch tools.
+
+[![DOI](https://zenodo.org/badge/472450059.svg)](https://zenodo.org/badge/latestdoi/472450059)
+![CI Workflow](https://github.com/nasa-nccs-hpda/pytorch-caney/actions/workflows/ci.yml/badge.svg)
+![CI to DockerHub ](https://github.com/nasa-nccs-hpda/pytorch-caney/actions/workflows/dockerhub.yml/badge.svg)
+![Code style: PEP8](https://github.com/nasa-nccs-hpda/pytorch-caney/actions/workflows/lint.yml/badge.svg)
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Coverage Status](https://coveralls.io/repos/github/nasa-nccs-hpda/pytorch-caney/badge.svg?branch=main)](https://coveralls.io/github/nasa-nccs-hpda/pytorch-caney?branch=main)
+
+## Documentation
+
+- Latest: https://nasa-nccs-hpda.github.io/pytorch-caney/latest
+
+## Objectives
+
+- Library to process remote sensing imagery using GPU and CPU parallelization.
+- Machine Learning and Deep Learning image classification and regression.
+- Agnostic array and vector-like data structures.
+- User interface environments via Notebooks for easy to use AI/ML projects.
+- Example notebooks for quick AI/ML start with your own data.
+
+## Installation
+
+The following library is intended to be used to accelerate the development of data science products
+for remote sensing satellite imagery, or any other applications. pytorch-caney can be installed
+by itself, but instructions for installing the full environments are listed under the requirements
+directory so projects, examples, and notebooks can be run.
+
+Note: PIP installations do not include CUDA libraries for GPU support. Make sure NVIDIA libraries
+are installed locally in the system if not using conda/mamba.
+
+```bash
+module load singularity # if a module needs to be loaded
+singularity build --sandbox pytorch-caney-container docker://nasanccs/pytorch-caney:latest
+```
+
+## Why Caney?
+
+"Caney" means longhouse in Taíno.
+
+## Contributors
+
+- Jordan Alexis Caraballo-Vega, jordan.a.caraballo-vega@nasa.gov
+- Caleb Spradlin, caleb.s.spradlin@nasa.gov
+
+## Contributing
+
+Please see our [guide for contributing to pytorch-caney](CONTRIBUTING.md).
+
+## SatVision
+
+| name | pretrain | resolution | #params |
+| :---: | :---: | :---: | :---: |
+| SatVision-B | MODIS-1.9-M | 192x192 | 84.5M |
+
+## SatVision Datasets
+
+| name | bands | resolution | #chips |
+| :---: | :---: | :---: | :---: |
+| MODIS-Small | 7 | 128x128 | 1,994,131 |
+
+## MODIS Surface Reflectance (MOD09GA) Band Details
+
+| Band Name      | Bandwidth     |
+| :------------: | :-----------: |
+| sur_refl_b01_1 | 0.620 - 0.670 |
+| sur_refl_b02_1 | 0.841 - 0.876 |
+| sur_refl_b03_1 | 0.459 - 0.479 |
+| sur_refl_b04_1 | 0.545 - 0.565 |
+| sur_refl_b05_1 | 1.230 - 1.250 |
+| sur_refl_b06_1 | 1.628 - 1.652 |
+| sur_refl_b07_1 | 2.105 - 2.155 |
+
+## Pre-training with Masked Image Modeling
+
+To pre-train the swinv2 base model with masked image modeling pre-training, run:
+```bash
+torchrun --nproc_per_node <NGPUS> pytorch-caney/pytorch_caney/pipelines/pretraining/mim.py --cfg <config-file> --dataset <dataset-name> --data-paths <path-to-data-subfolder-1> --batch-size <batch-size> --output <output-dir> --enable-amp
+```
+
+For example to run on a compute node with 4 GPUs and a batch size of 128 on the MODIS SatVision pre-training dataset with a base swinv2 model, run:
+
+```bash
+singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+Singularity> export PYTHONPATH=$PWD:$PWD/pytorch-caney
+Singularity> torchrun --nproc_per_node 4 pytorch-caney/pytorch_caney/pipelines/pretraining/mim.py --cfg pytorch-caney/examples/satvision/mim_pretrain_swinv2_satvision_base_192_window12_800ep.yaml --dataset MODIS --data-paths /explore/nobackup/projects/ilab/data/satvision/pretraining/training_* --batch-size 128 --output . --enable-amp
+```
+
+This example script runs the exact configuration used to make the SatVision-base model pre-training with MiM and the MODIS pre-training dataset.
+```bash
+singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+Singularity> cd pytorch-caney/examples/satvision
+Singularity> ./run_satvision_pretrain.sh
+```
+
+## Fine-tuning Satvision-base
+To fine-tune the satvision-base pre-trained model, run:
+```bash
+torchrun --nproc_per_node <NGPUS> pytorch-caney/pytorch_caney/pipelines/finetuning/finetune.py --cfg <config-file> --pretrained <path-to-pretrained> --dataset <dataset-name> --data-paths <path-to-data-subfolder-1> --batch-size <batch-size> --output <output-dir> --enable-amp
+```
+
+See example config files pytorch-caney/examples/satvision/finetune_satvision_base_*.yaml to see how to structure your config file for fine-tuning.
+
+
+## Testing
+For unittests, run this bash command to run linting and unit test runs. This will execute unit tests and linting in a temporary venv environment only used for testing.
+```bash
+git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+cd pytorch-caney; bash test.sh
+```
+or run unit tests directly with container or anaconda env
+
+```bash
+git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+singularity build --sandbox pytorch-caney-container docker://nasanccs/pytorch-caney:latest
+singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+cd pytorch-caney; python -m unittest discover pytorch_caney/tests
+```
+
+```bash
+git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+cd pytorch-caney; conda env create -f requirements/environment_gpu.yml;
+conda activate pytorch-caney
+python -m unittest discover pytorch_caney/tests
+```
+## References
+
+- [Pytorch Lightning](https://github.com/Lightning-AI/lightning)
+- [Swin Transformer](https://github.com/microsoft/Swin-Transformer)
+- [SimMIM](https://github.com/microsoft/SimMIM)

pytorch-caney/README.rst

@@ -0,0 +1,164 @@
+================
+pytorch-caney
+================
+
+Python package for lots of Pytorch tools for geospatial science problems.
+
+.. image:: https://zenodo.org/badge/472450059.svg
+      :target: https://zenodo.org/badge/latestdoi/472450059
+
+Objectives
+------------
+
+- Library to process remote sensing imagery using GPU and CPU parallelization.
+- Machine Learning and Deep Learning image classification and regression.
+- Agnostic array and vector-like data structures.
+- User interface environments via Notebooks for easy to use AI/ML projects.
+- Example notebooks for quick AI/ML start with your own data.
+
+Installation
+----------------
+
+The following library is intended to be used to accelerate the development of data science products
+for remote sensing satellite imagery, or any other applications. pytorch-caney can be installed
+by itself, but instructions for installing the full environments are listed under the requirements
+directory so projects, examples, and notebooks can be run.
+
+Note: PIP installations do not include CUDA libraries for GPU support. Make sure NVIDIA libraries
+are installed locally in the system if not using conda/mamba.
+
+.. code-block:: bash
+
+    module load singularity # if a module needs to be loaded
+    singularity build --sandbox pytorch-caney-container docker://nasanccs/pytorch-caney:latest
+
+
+Why Caney?
+---------------
+
+"Caney" means longhouse in Taíno.
+
+Contributors
+-------------
+
+- Jordan Alexis Caraballo-Vega, jordan.a.caraballo-vega@nasa.gov
+- Caleb Spradlin, caleb.s.spradlin@nasa.gov
+- Jian Li, jian.li@nasa.gov
+
+Contributing
+-------------
+
+Please see our `guide for contributing to pytorch-caney <CONTRIBUTING.md>`_.
+
+SatVision
+------------
+
++---------------+--------------+------------+------------+
+| Name          | Pretrain     | Resolution | Parameters |
++===============+==============+============+============+
+| SatVision-B   | MODIS-1.9-M  | 192x192    | 84.5M      |
++---------------+--------------+------------+------------+
+
+SatVision Datasets
+-----------------------
+
++---------------+-----------+------------+-------------+
+| Name          | Bands     | Resolution | Image Chips |
++===============+===========+============+=============+
+| MODIS-Small   | 7         | 128x128    | 1,994,131   |
++---------------+-----------+------------+-------------+
+
+MODIS Surface Reflectance (MOD09GA) Band Details
+------------------------------------------------------
+
++-----------------+---------------+
+| Band Name       | Bandwidth     |
++=================+===============+
+| sur_refl_b01_1  | 0.620 - 0.670 |
++-----------------+---------------+
+| sur_refl_b02_1  | 0.841 - 0.876 |
++-----------------+---------------+
+| sur_refl_b03_1  | 0.459 - 0.479 |
++-----------------+---------------+
+| sur_refl_b04_1  | 0.545 - 0.565 |
++-----------------+---------------+
+| sur_refl_b05_1  | 1.230 - 1.250 |
++-----------------+---------------+
+| sur_refl_b06_1  | 1.628 - 1.652 |
++-----------------+---------------+
+| sur_refl_b07_1  | 2.105 - 2.155 |
++-----------------+---------------+
+
+Pre-training with Masked Image Modeling
+-----------------------------------------
+
+To pre-train the swinv2 base model with masked image modeling pre-training, run:
+
+.. code-block:: bash
+
+    torchrun --nproc_per_node <NGPUS> pytorch-caney/pytorch_caney/pipelines/pretraining/mim.py --cfg <config-file> --dataset <dataset-name> --data-paths <path-to-data-subfolder-1> --batch-size <batch-size> --output <output-dir> --enable-amp
+
+For example to run on a compute node with 4 GPUs and a batch size of 128 on the MODIS SatVision pre-training dataset with a base swinv2 model, run:
+
+.. code-block:: bash
+
+    singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+    Singularity> export PYTHONPATH=$PWD:$PWD/pytorch-caney
+    Singularity> torchrun --nproc_per_node 4 pytorch-caney/pytorch_caney/pipelines/pretraining/mim.py --cfg pytorch-caney/examples/satvision/mim_pretrain_swinv2_satvision_base_192_window12_800ep.yaml --dataset MODIS --data-paths /explore/nobackup/projects/ilab/data/satvision/pretraining/training_* --batch-size 128 --output . --enable-amp
+
+
+This example script runs the exact configuration used to make the SatVision-base model pre-training with MiM and the MODIS pre-training dataset.
+
+.. code-block:: bash
+
+    singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+    Singularity> cd pytorch-caney/examples/satvision
+    Singularity> ./run_satvision_pretrain.sh
+
+
+Fine-tuning Satvision-base
+-----------------------------
+
+To fine-tune the satvision-base pre-trained model, run:
+
+.. code-block:: bash
+
+    torchrun --nproc_per_node <NGPUS> pytorch-caney/pytorch_caney/pipelines/finetuning/finetune.py --cfg <config-file> --pretrained <path-to-pretrained> --dataset <dataset-name> --data-paths <path-to-data-subfolder-1> --batch-size <batch-size> --output <output-dir> --enable-amp
+
+See example config files pytorch-caney/examples/satvision/finetune_satvision_base_*.yaml to see how to structure your config file for fine-tuning.
+
+
+Testing
+------------
+
+For unittests, run this bash command to run linting and unit test runs. This will execute unit tests and linting in a temporary venv environment only used for testing.
+
+.. code-block:: bash
+
+    git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+    cd pytorch-caney; bash test.sh
+
+
+or run unit tests directly with container or anaconda env
+
+.. code-block:: bash
+
+    git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+    singularity build --sandbox pytorch-caney-container docker://nasanccs/pytorch-caney:latest
+    singularity shell --nv -B <mounts> /path/to/container/pytorch-caney-container
+    cd pytorch-caney; python -m unittest discover pytorch_caney/tests
+
+.. code-block:: bash
+
+    git clone git@github.com:nasa-nccs-hpda/pytorch-caney.git
+    cd pytorch-caney; conda env create -f requirements/environment_gpu.yml;
+    conda activate pytorch-caney
+    python -m unittest discover pytorch_caney/tests
+
+
+References
+------------
+
+- `Pytorch Lightning <https://github.com/Lightning-AI/lightning>`_ 
+- `Swin Transformer <https://github.com/microsoft/Swin-Transformer>`_ 
+- `SimMIM <https://github.com/microsoft/SimMIM>`_ 

pytorch-caney/docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

pytorch-caney/docs/conf.py

@@ -0,0 +1,52 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath('..'))
+
+import pytorch_caney  # noqa: E402
+
+project = 'pytorch-caney'
+copyright = '2023, Jordan A. Caraballo-Vega'
+author = 'Jordan A. Caraballo-Vega'
+
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx_autodoc_typehints',
+    'jupyter_sphinx.execute',
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.viewcode",
+    "sphinx.ext.napoleon",
+    "sphinx_click.ext",
+    "sphinx.ext.githubpages",
+    "nbsphinx",
+]
+
+intersphinx_mapping = {
+    "pyproj": ("https://pyproj4.github.io/pyproj/stable/", None),
+    "rasterio": ("https://rasterio.readthedocs.io/en/stable/", None),
+    "xarray": ("http://xarray.pydata.org/en/stable/", None),
+}
+
+templates_path = ['_templates']
+exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+
+master_doc = "index"
+
+version = release = pytorch_caney.__version__
+
+pygments_style = "sphinx"
+
+todo_include_todos = False
+
+html_theme = 'sphinx_rtd_theme'
+html_logo = 'static/DSG_LOGO_REDESIGN.png'
+
+myst_enable_extensions = [
+    "amsmath",
+    "colon_fence",
+    "deflist",
+    "dollarmath",
+    "html_image",
+]
+
+myst_url_schemes = ("http", "https", "mailto")

pytorch-caney/docs/examples.rst

@@ -0,0 +1,3 @@
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:

pytorch-caney/docs/index.rst

@@ -0,0 +1,22 @@
+.. pytorch-caney's documentation master file, created by
+   sphinx-quickstart on Fri Jun 23 11:32:18 2023.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+Welcome to pytorch-caney's documentation!
+=========================================
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   readme
+   examples
+   modules
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

pytorch-caney/docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

pytorch-caney/docs/modules.rst

@@ -0,0 +1,284 @@
+pytorch-caney package
+========================
+
+pytorch_caney.config
+----------------------
+
+.. automodule:: pytorch_caney.config
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datamodules.finetune_datamodule
+----------------------
+
+.. automodule:: pytorch_caney.data.datamodules.finetune_datamodule
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datamodules.mim_datamodule
+----------------------
+
+.. automodule:: pytorch_caney.data.datamodules.mim_datamodule
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datamodules.segmentation_datamodule
+----------------------
+
+.. automodule:: pytorch_caney.data.datamodules.segmentation_datamodule
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datamodules.simmim_datamodule
+----------------------
+
+.. automodule:: pytorch_caney.data.datamodules.simmim_datamodule
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datasets.classification_dataset
+----------------------
+
+.. automodule:: pytorch_caney.data.datasets.classification_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.datasets.modis_dataset
+----------------------
+
+.. automodule:: pytorch_caney.data.datasets.modis_dataset
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.transforms
+----------------------
+
+.. automodule:: pytorch_caney.data.transforms
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.data.utils
+----------------------
+
+.. automodule:: pytorch_caney.data.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.inference
+----------------------
+
+.. automodule:: pytorch_caney.inference
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.loss.build
+----------------------
+
+.. automodule:: pytorch_caney.loss.build
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.loss.utils
+----------------------
+
+.. automodule:: pytorch_caney.loss.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.lr_scheduler
+----------------------
+
+.. automodule:: pytorch_caney.lr_scheduler
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.metrics
+----------------------
+
+.. automodule:: pytorch_caney.metrics
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.decoders.unet_decoder
+----------------------
+
+.. automodule:: pytorch_caney.models.decoders.unet_decoder
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.mim.mim
+----------------------
+
+.. automodule:: pytorch_caney.models.mim.mim
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.simmim.simmim
+----------------------
+
+.. automodule:: pytorch_caney.models.simmim.simmim
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.build
+----------------------
+
+.. automodule:: pytorch_caney.models.build
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.maskrcnn_model
+----------------------
+
+.. automodule:: pytorch_caney.models.maskrcnn_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.swinv2_model
+----------------------
+
+.. automodule:: pytorch_caney.models.swinv2_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.unet_model
+----------------------
+
+.. automodule:: pytorch_caney.models.unet_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.models.unet_swin_model
+----------------------
+
+.. automodule:: pytorch_caney.models.unet_swin_model
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.network.attention
+----------------------
+
+.. automodule:: pytorch_caney.network.attention
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.network.mlp
+----------------------
+
+.. automodule:: pytorch_caney.network.mlp
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.pipelines.finetuning.finetune
+----------------------
+
+.. automodule:: pytorch_caney.pipelines.finetuning.finetune
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.pipelines.pretraining.mim
+----------------------
+
+.. automodule:: pytorch_caney.pipelines.pretraining.mim
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.pipelines.modis_segmentation
+----------------------
+
+.. automodule:: pytorch_caney.pipelines.modis_segmentation
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.processing
+----------------------
+
+.. automodule:: pytorch_caney.processing
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.ptc_logging
+----------------------
+
+.. automodule:: pytorch_caney.ptc_logging
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.training.fine_tuning
+----------------------
+
+.. automodule:: pytorch_caney.training.fine_tuning
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.training.mim_utils
+----------------------
+
+.. automodule:: pytorch_caney.training.mim_utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.training.pre_training
+----------------------
+
+.. automodule:: pytorch_caney.training.pre_training
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.training.simmim_utils
+----------------------
+
+.. automodule:: pytorch_caney.training.simmim_utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.training.utils
+----------------------
+
+.. automodule:: pytorch_caney.training.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+pytorch_caney.utils
+----------------------
+
+.. automodule:: pytorch_caney.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+

pytorch-caney/docs/pytorch_caney.rst

@@ -0,0 +1,31 @@
+pytorch-caney
+------------------
+
+Python package for lots of Pytorch tools for geospatial science problems.
+
+Installation
+-----------------
+
+Install with pip
+::
+
+    pip install pytorch-caney
+
+API
+----
+
+.. automodule:: pytorch_caney
+   :members:
+
+Authors
+----------
+
+Jordan A. Caraballo-Vega, jordan.a.caraballo-vega@nasa.gov
+Caleb S. Spradlin, caleb.s.spradlin@nasa.gov
+Jian Li, jian.li@nasa.gov
+
+License
+---------------
+
+The package is released under the `MIT
+License <https://opensource.org/licenses/MIT>`__.

pytorch-caney/docs/readme.rst

@@ -0,0 +1 @@
+.. include:: ../README.rst

pytorch-caney/docs/requirements.txt

@@ -0,0 +1,4 @@
+sphinx
+sphinx-autodoc-typehints
+jupyter-sphinx
+sphinx-rtd-theme

pytorch-caney/docs/source/index.rst

@@ -0,0 +1,20 @@
+.. tensorflow-caney documentation master file, created by
+   sphinx-quickstart on Fri Jan 13 06:59:19 2023.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+Welcome to tensorflow-caney's documentation!
+============================================
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

pytorch-caney/examples/satvision/finetune_satvision_base_landcover5class_192_window12_100ep.yaml

@@ -0,0 +1,33 @@
+MODEL:
+  TYPE: swinv2
+  DECODER: unet
+  NAME: satvision_finetune_lc5class
+  DROP_PATH_RATE: 0.1
+  NUM_CLASSES: 5
+  SWINV2:
+    IN_CHANS: 7
+    EMBED_DIM: 128
+    DEPTHS: [ 2, 2, 18, 2 ]
+    NUM_HEADS: [ 4, 8, 16, 32 ]
+    WINDOW_SIZE: 14
+    PRETRAINED_WINDOW_SIZES: [ 12, 12, 12, 6 ]
+DATA:
+  IMG_SIZE: 224
+  DATASET: MODISLC5
+  MASK_PATCH_SIZE: 32
+  MASK_RATIO: 0.6
+LOSS:
+  NAME: 'tversky'
+  MODE: 'multiclass'
+  ALPHA: 0.4
+  BETA: 0.6
+TRAIN:
+  EPOCHS: 100
+  WARMUP_EPOCHS: 10
+  BASE_LR: 1e-4
+  WARMUP_LR: 5e-7
+  WEIGHT_DECAY: 0.01
+  LAYER_DECAY: 0.8
+PRINT_FREQ: 100
+SAVE_FREQ: 5
+TAG: satvision_finetune_land_cover_5class_swinv2_satvision_192_window12__800ep

pytorch-caney/examples/satvision/finetune_satvision_base_landcover9class_192_window12_100ep.yaml

@@ -0,0 +1,33 @@
+MODEL:
+  TYPE: swinv2
+  DECODER: unet
+  NAME: satvision_finetune_lc9class
+  DROP_PATH_RATE: 0.1
+  NUM_CLASSES: 9
+  SWINV2:
+    IN_CHANS: 7
+    EMBED_DIM: 128
+    DEPTHS: [ 2, 2, 18, 2 ]
+    NUM_HEADS: [ 4, 8, 16, 32 ]
+    WINDOW_SIZE: 14
+    PRETRAINED_WINDOW_SIZES: [ 12, 12, 12, 6 ]
+DATA:
+  IMG_SIZE: 224
+  DATASET: MODISLC5
+  MASK_PATCH_SIZE: 32
+  MASK_RATIO: 0.6
+LOSS:
+  NAME: 'tversky'
+  MODE: 'multiclass'
+  ALPHA: 0.4
+  BETA: 0.6
+TRAIN:
+  EPOCHS: 100
+  WARMUP_EPOCHS: 10
+  BASE_LR: 1e-4
+  WARMUP_LR: 5e-7
+  WEIGHT_DECAY: 0.01
+  LAYER_DECAY: 0.8
+PRINT_FREQ: 100
+SAVE_FREQ: 5
+TAG: satvision_finetune_land_cover_9class_swinv2_satvision_192_window12__800ep

pytorch-caney/examples/satvision/mim_pretrain_swinv2_satvision_base_192_window12_800ep.yaml

@@ -0,0 +1,27 @@
+MODEL:
+  TYPE: swinv2
+  NAME: mim_satvision_pretrain
+  DROP_PATH_RATE: 0.1
+  SWINV2:
+    IN_CHANS: 7
+    EMBED_DIM: 128
+    DEPTHS: [ 2, 2, 18, 2 ]
+    NUM_HEADS: [ 4, 8, 16, 32 ]
+    WINDOW_SIZE: 12
+DATA:
+  IMG_SIZE: 192
+  MASK_PATCH_SIZE: 32
+  MASK_RATIO: 0.6
+TRAIN:
+  EPOCHS: 800
+  WARMUP_EPOCHS: 10
+  BASE_LR: 1e-4
+  WARMUP_LR: 5e-7
+  WEIGHT_DECAY: 0.05
+  LR_SCHEDULER:
+    NAME: 'multistep'
+    GAMMA: 0.1
+    MULTISTEPS: [700,]
+PRINT_FREQ: 100
+SAVE_FREQ: 5
+TAG: mim_pretrain_swinv2_satvision_192_window12__800ep

pytorch-caney/examples/satvision/run_satvision_finetune_lc_fiveclass.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+#SBATCH -J finetune_satvision_lc5
+#SBATCH -t 3-00:00:00
+#SBATCH -G 4
+#SBATCH -N 1
+
+
+export PYTHONPATH=$PWD:../../../:../../../pytorch-caney
+export NGPUS=8
+
+torchrun --nproc_per_node $NGPUS \
+    ../../../pytorch-caney/pytorch_caney/pipelines/finetuning/finetune.py \
+	--cfg finetune_satvision_base_landcover5class_192_window12_100ep.yaml \
+	--pretrained /explore/nobackup/people/cssprad1/projects/satnet/code/development/masked_image_modeling/development/models/simmim_satnet_pretrain_pretrain/simmim_pretrain__satnet_swinv2_base__img192_window12__800ep_v3_no_norm/ckpt_epoch_800.pth \
+    --dataset MODISLC9 \
+	--data-paths /explore/nobackup/projects/ilab/data/satvision/finetuning/h18v04/labels_9classes_224 \
+	--batch-size 4 \
+	--output /explore/nobackup/people/cssprad1/projects/satnet/code/development/cleanup/finetune/models \
+	--enable-amp

pytorch-caney/examples/satvision/run_satvision_finetune_lc_nineclass.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+#SBATCH -J finetune_satvision_lc9
+#SBATCH -t 3-00:00:00
+#SBATCH -G 4
+#SBATCH -N 1
+
+
+export PYTHONPATH=$PWD:../../../:../../../pytorch-caney
+export NGPUS=8
+
+torchrun --nproc_per_node $NGPUS \
+    ../../../pytorch-caney/pytorch_caney/pipelines/finetuning/finetune.py \
+	--cfg finetune_satvision_base_landcover5class_192_window12_100ep.yaml \
+	--pretrained /explore/nobackup/people/cssprad1/projects/satnet/code/development/masked_image_modeling/development/models/simmim_satnet_pretrain_pretrain/simmim_pretrain__satnet_swinv2_base__img192_window12__800ep_v3_no_norm/ckpt_epoch_800.pth \
+    --dataset MODISLC9 \
+	--data-paths /explore/nobackup/projects/ilab/data/satvision/finetuning/h18v04/labels_5classes_224 \
+	--batch-size 4 \
+	--output /explore/nobackup/people/cssprad1/projects/satnet/code/development/cleanup/finetune/models \
+	--enable-amp

pytorch-caney/examples/satvision/run_satvision_pretrain.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+#SBATCH -J pretrain_satvision_swinv2
+#SBATCH -t 3-00:00:00
+#SBATCH -G 4
+#SBATCH -N 1
+
+
+export PYTHONPATH=$PWD:../../../:../../../pytorch-caney
+export NGPUS=4
+
+torchrun --nproc_per_node $NGPUS \
+    ../../../pytorch-caney/pytorch_caney/pipelines/pretraining/mim.py \
+	--cfg mim_pretrain_swinv2_satvision_base_192_window12_800ep.yaml \
+    --dataset MODIS \
+	--data-paths /explore/nobackup/projects/ilab/data/satvision/pretraining/training_* \
+	--batch-size 128 \
+	--output /explore/nobackup/people/cssprad1/projects/satnet/code/development/cleanup/trf/transformer/models \
+	--enable-amp

pytorch-caney/pyproject.toml

@@ -0,0 +1,6 @@
+[build-system]
+# Minimum requirements for the build system to execute.
+requires = ["setuptools", "wheel"]
+
+[tool.black]
+target_version = ['py39']

pytorch-caney/pytorch_caney/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pytorch-caney/pytorch_caney/config.py

@@ -0,0 +1,226 @@
+import os
+import yaml
+from yacs.config import CfgNode as CN
+
+_C = CN()
+
+# Base config files
+_C.BASE = ['']
+
+# -----------------------------------------------------------------------------
+# Data settings
+# -----------------------------------------------------------------------------
+_C.DATA = CN()
+# Batch size for a single GPU, could be overwritten by command line argument
+_C.DATA.BATCH_SIZE = 128
+# Path(s) to dataset, could be overwritten by command line argument
+_C.DATA.DATA_PATHS = ['']
+# Dataset name
+_C.DATA.DATASET = 'MODIS'
+# Input image size
+_C.DATA.IMG_SIZE = 224
+# Interpolation to resize image (random, bilinear, bicubic)
+_C.DATA.INTERPOLATION = 'bicubic'
+# Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.
+_C.DATA.PIN_MEMORY = True
+# Number of data loading threads
+_C.DATA.NUM_WORKERS = 8
+# [SimMIM] Mask patch size for MaskGenerator
+_C.DATA.MASK_PATCH_SIZE = 32
+# [SimMIM] Mask ratio for MaskGenerator
+_C.DATA.MASK_RATIO = 0.6
+
+# -----------------------------------------------------------------------------
+# Model settings
+# -----------------------------------------------------------------------------
+_C.MODEL = CN()
+# Model type
+_C.MODEL.TYPE = 'swinv2'
+# Decoder type
+_C.MODEL.DECODER = None
+# Model name
+_C.MODEL.NAME = 'swinv2_base_patch4_window7_224'
+# Pretrained weight from checkpoint, could be from previous pre-training
+# could be overwritten by command line argument
+_C.MODEL.PRETRAINED = ''
+# Checkpoint to resume, could be overwritten by command line argument
+_C.MODEL.RESUME = ''
+# Number of classes, overwritten in data preparation
+_C.MODEL.NUM_CLASSES = 17
+# Dropout rate
+_C.MODEL.DROP_RATE = 0.0
+# Drop path rate
+_C.MODEL.DROP_PATH_RATE = 0.1
+
+# Swin Transformer V2 parameters
+_C.MODEL.SWINV2 = CN()
+_C.MODEL.SWINV2.PATCH_SIZE = 4
+_C.MODEL.SWINV2.IN_CHANS = 3
+_C.MODEL.SWINV2.EMBED_DIM = 96
+_C.MODEL.SWINV2.DEPTHS = [2, 2, 6, 2]
+_C.MODEL.SWINV2.NUM_HEADS = [3, 6, 12, 24]
+_C.MODEL.SWINV2.WINDOW_SIZE = 7
+_C.MODEL.SWINV2.MLP_RATIO = 4.
+_C.MODEL.SWINV2.QKV_BIAS = True
+_C.MODEL.SWINV2.APE = False
+_C.MODEL.SWINV2.PATCH_NORM = True
+_C.MODEL.SWINV2.PRETRAINED_WINDOW_SIZES = [0, 0, 0, 0]
+
+# -----------------------------------------------------------------------------
+# Training settings
+# -----------------------------------------------------------------------------
+_C.LOSS = CN()
+_C.LOSS.NAME = 'tversky'
+_C.LOSS.MODE = 'multiclass'
+_C.LOSS.CLASSES = None
+_C.LOSS.LOG = False
+_C.LOSS.LOGITS = True
+_C.LOSS.SMOOTH = 0.0
+_C.LOSS.IGNORE_INDEX = None
+_C.LOSS.EPS = 1e-7
+_C.LOSS.ALPHA = 0.5
+_C.LOSS.BETA = 0.5
+_C.LOSS.GAMMA = 1.0
+
+# -----------------------------------------------------------------------------
+# Training settings
+# -----------------------------------------------------------------------------
+_C.TRAIN = CN()
+_C.TRAIN.START_EPOCH = 0
+_C.TRAIN.EPOCHS = 300
+_C.TRAIN.WARMUP_EPOCHS = 20
+_C.TRAIN.WEIGHT_DECAY = 0.05
+_C.TRAIN.BASE_LR = 5e-4
+_C.TRAIN.WARMUP_LR = 5e-7
+_C.TRAIN.MIN_LR = 5e-6
+# Clip gradient norm
+_C.TRAIN.CLIP_GRAD = 5.0
+# Auto resume from latest checkpoint
+_C.TRAIN.AUTO_RESUME = True
+# Gradient accumulation steps
+# could be overwritten by command line argument
+_C.TRAIN.ACCUMULATION_STEPS = 0
+# Whether to use gradient checkpointing to save memory
+# could be overwritten by command line argument
+_C.TRAIN.USE_CHECKPOINT = False
+
+# LR scheduler
+_C.TRAIN.LR_SCHEDULER = CN()
+_C.TRAIN.LR_SCHEDULER.NAME = 'cosine'
+# Epoch interval to decay LR, used in StepLRScheduler
+_C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30
+# LR decay rate, used in StepLRScheduler
+_C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1
+# Gamma / Multi steps value, used in MultiStepLRScheduler
+_C.TRAIN.LR_SCHEDULER.GAMMA = 0.1
+_C.TRAIN.LR_SCHEDULER.MULTISTEPS = []
+
+# Optimizer
+_C.TRAIN.OPTIMIZER = CN()
+_C.TRAIN.OPTIMIZER.NAME = 'adamw'
+# Optimizer Epsilon
+_C.TRAIN.OPTIMIZER.EPS = 1e-8
+# Optimizer Betas
+_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999)
+# SGD momentum
+_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9
+
+# [SimMIM] Layer decay for fine-tuning
+_C.TRAIN.LAYER_DECAY = 1.0
+
+
+# -----------------------------------------------------------------------------
+# Testing settings
+# -----------------------------------------------------------------------------
+_C.TEST = CN()
+# Whether to use center crop when testing
+_C.TEST.CROP = True
+
+# -----------------------------------------------------------------------------
+# Misc
+# -----------------------------------------------------------------------------
+# Whether to enable pytorch amp, overwritten by command line argument
+_C.ENABLE_AMP = False
+# Enable Pytorch automatic mixed precision (amp).
+_C.AMP_ENABLE = True
+# Path to output folder, overwritten by command line argument
+_C.OUTPUT = ''
+# Tag of experiment, overwritten by command line argument
+_C.TAG = 'pt-caney-default-tag'
+# Frequency to save checkpoint
+_C.SAVE_FREQ = 1
+# Frequency to logging info
+_C.PRINT_FREQ = 10
+# Fixed random seed
+_C.SEED = 42
+# Perform evaluation only, overwritten by command line argument
+_C.EVAL_MODE = False
+
+
+def _update_config_from_file(config, cfg_file):
+    config.defrost()
+    with open(cfg_file, 'r') as f:
+        yaml_cfg = yaml.load(f, Loader=yaml.FullLoader)
+
+    for cfg in yaml_cfg.setdefault('BASE', ['']):
+        if cfg:
+            _update_config_from_file(
+                config, os.path.join(os.path.dirname(cfg_file), cfg)
+            )
+    print('=> merge config from {}'.format(cfg_file))
+    config.merge_from_file(cfg_file)
+    config.freeze()
+
+
+def update_config(config, args):
+    _update_config_from_file(config, args.cfg)
+
+    config.defrost()
+
+    def _check_args(name):
+        if hasattr(args, name) and eval(f'args.{name}'):
+            return True
+        return False
+
+    # merge from specific arguments
+    if _check_args('batch_size'):
+        config.DATA.BATCH_SIZE = args.batch_size
+    if _check_args('data_paths'):
+        config.DATA.DATA_PATHS = args.data_paths
+    if _check_args('dataset'):
+        config.DATA.DATASET = args.dataset
+    if _check_args('resume'):
+        config.MODEL.RESUME = args.resume
+    if _check_args('pretrained'):
+        config.MODEL.PRETRAINED = args.pretrained
+    if _check_args('resume'):
+        config.MODEL.RESUME = args.resume
+    if _check_args('accumulation_steps'):
+        config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps
+    if _check_args('use_checkpoint'):
+        config.TRAIN.USE_CHECKPOINT = True
+    if _check_args('disable_amp'):
+        config.AMP_ENABLE = False
+    if _check_args('output'):
+        config.OUTPUT = args.output
+    if _check_args('tag'):
+        config.TAG = args.tag
+    if _check_args('eval'):
+        config.EVAL_MODE = True
+    if _check_args('enable_amp'):
+        config.ENABLE_AMP = args.enable_amp
+
+    # output folder
+    config.OUTPUT = os.path.join(config.OUTPUT, config.MODEL.NAME, config.TAG)
+
+    config.freeze()
+
+
+def get_config(args):
+    """Get a yacs CfgNode object with default values."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    config = _C.clone()
+    update_config(config, args)
+
+    return config

pytorch-caney/pytorch_caney/console/cli.py

@@ -0,0 +1,62 @@
+from pytorch_lightning.utilities.cli import LightningCLI
+
+import torch
+
+
+class TerraGPULightningCLI(LightningCLI):
+
+    def add_arguments_to_parser(self, parser):
+
+        # Trainer - performance
+        parser.set_defaults({"trainer.accelerator": "auto"})
+        parser.set_defaults({"trainer.devices": "auto"})
+        parser.set_defaults({"trainer.auto_select_gpus": True})
+        parser.set_defaults({"trainer.precision": 32})
+
+        # Trainer - training
+        parser.set_defaults({"trainer.max_epochs": 500})
+        parser.set_defaults({"trainer.min_epochs": 1})
+        parser.set_defaults({"trainer.detect_anomaly": True})
+        parser.set_defaults({"trainer.logger": True})
+        parser.set_defaults({"trainer.default_root_dir": "output_model"})
+
+        # Trainer - optimizer - TODO
+        _ = {
+            "class_path": torch.optim.Adam,
+            "init_args": {
+                "lr": 0.01
+            }
+        }
+
+        # Trainer - callbacks
+        default_callbacks = [
+            {"class_path": "pytorch_lightning.callbacks.DeviceStatsMonitor"},
+            {
+                "class_path": "pytorch_lightning.callbacks.EarlyStopping",
+                "init_args": {
+                    "monitor": "val_loss",
+                    "patience": 5,
+                    "mode": "min"
+                }
+            },
+            # {
+            #    "class_path": "pytorch_lightning.callbacks.ModelCheckpoint",
+            #    "init_args": {
+            #        "dirpath": "output_model",
+            #        "monitor": "val_loss",
+            #        "auto_insert_metric_name": True
+            #    }
+            # },
+        ]
+        parser.set_defaults({"trainer.callbacks": default_callbacks})
+
+        #    {
+        #        "class_path": "pytorch_lightning.callbacks.ModelCheckpoint",
+        #        "init_args": {
+        #            "dirpath": "output_model",
+        #            "monitor": "val_loss",
+        #            "auto_insert_metric_name": True
+        #        }
+        #    },
+        # ]
+        # parser.set_defaults({"trainer.callbacks": default_callbacks})

pytorch-caney/pytorch_caney/console/dl_pipeline.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+# RF pipeline: preprocess, train, and predict.
+
+import sys
+import logging
+
+# from terragpu import unet_model
+# from terragpu.decorators import DuplicateFilter
+# from terragpu.ai.deep_learning.datamodules.segmentation_datamodule \
+# import SegmentationDataModule
+
+from pytorch_lightning import seed_everything  # , trainer
+# from pytorch_lightning import LightningModule, LightningDataModule
+from terragpu.ai.deep_learning.console.cli import TerraGPULightningCLI
+
+
+# -----------------------------------------------------------------------------
+# main
+#
+# python rf_pipeline.py options here
+# -----------------------------------------------------------------------------
+def main():
+
+    # -------------------------------------------------------------------------
+    # Set logging
+    # -------------------------------------------------------------------------
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    ch = logging.StreamHandler(sys.stdout)
+    ch.setLevel(logging.INFO)
+
+    # Set formatter and handlers
+    formatter = logging.Formatter(
+        "%(asctime)s; %(levelname)s; %(message)s", "%Y-%m-%d %H:%M:%S")
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+
+    # -------------------------------------------------------------------------
+    # Execute pipeline step
+    # -------------------------------------------------------------------------
+    # Seed every library
+    seed_everything(1234, workers=True)
+    _ = TerraGPULightningCLI(save_config_callback=None)
+    # unet_model.UNetSegmentation, SegmentationDataModule)
+
+    # train
+    # trainer = pl.Trainer()
+    # trainer.fit(model, datamodule=dm)
+    # validate
+    # trainer.validate(datamodule=dm)
+    # test
+    # trainer.test(datamodule=dm)
+    # predict
+    # predictions = trainer.predict(datamodule=dm)
+    return
+
+
+# -----------------------------------------------------------------------------
+# Invoke the main
+# -----------------------------------------------------------------------------
+if __name__ == "__main__":
+    sys.exit(main())

pytorch-caney/pytorch_caney/data/datamodules/finetune_datamodule.py

@@ -0,0 +1,114 @@
+from ..datasets.modis_dataset import MODISDataset
+from ..datasets.modis_lc_five_dataset import MODISLCFiveDataset
+from ..datasets.modis_lc_nine_dataset import MODISLCNineDataset
+
+from ..transforms import TensorResizeTransform
+
+import torch.distributed as dist
+from torch.utils.data import DataLoader, DistributedSampler
+
+
+DATASETS = {
+    'modis': MODISDataset,
+    'modislc9': MODISLCNineDataset,
+    'modislc5': MODISLCFiveDataset,
+    # 'modis tree': MODISTree,
+}
+
+
+def get_dataset_from_dict(dataset_name: str):
+    """Gets the proper dataset given a dataset name.
+
+    Args:
+        dataset_name (str): name of the dataset
+
+    Raises:
+        KeyError: thrown if dataset key is not present in dict
+
+    Returns:
+        dataset: pytorch dataset
+    """
+
+    dataset_name = dataset_name.lower()
+
+    try:
+
+        dataset_to_use = DATASETS[dataset_name]
+
+    except KeyError:
+
+        error_msg = f"{dataset_name} is not an existing dataset"
+
+        error_msg = f"{error_msg}. Available datasets: {DATASETS.keys()}"
+
+        raise KeyError(error_msg)
+
+    return dataset_to_use
+
+
+def build_finetune_dataloaders(config, logger):
+    """Builds the dataloaders and datasets for a fine-tuning task.
+
+    Args:
+        config: config object
+        logger: logging logger
+
+    Returns:
+        dataloader_train: training dataloader
+        dataloader_val: validation dataloader
+    """
+
+    transform = TensorResizeTransform(config)
+
+    logger.info(f'Finetuning data transform:\n{transform}')
+
+    dataset_name = config.DATA.DATASET
+
+    logger.info(f'Dataset: {dataset_name}')
+    logger.info(f'Data Paths: {config.DATA.DATA_PATHS}')
+
+    dataset_to_use = get_dataset_from_dict(dataset_name)
+
+    logger.info(f'Dataset obj: {dataset_to_use}')
+
+    dataset_train = dataset_to_use(data_paths=config.DATA.DATA_PATHS,
+                                   split="train",
+                                   img_size=config.DATA.IMG_SIZE,
+                                   transform=transform)
+
+    dataset_val = dataset_to_use(data_paths=config.DATA.DATA_PATHS,
+                                 split="val",
+                                 img_size=config.DATA.IMG_SIZE,
+                                 transform=transform)
+
+    logger.info(f'Build dataset: train images = {len(dataset_train)}')
+
+    logger.info(f'Build dataset: val images = {len(dataset_val)}')
+
+    sampler_train = DistributedSampler(
+        dataset_train,
+        num_replicas=dist.get_world_size(),
+        rank=dist.get_rank(),
+        shuffle=True)
+
+    sampler_val = DistributedSampler(
+        dataset_val,
+        num_replicas=dist.get_world_size(),
+        rank=dist.get_rank(),
+        shuffle=False)
+
+    dataloader_train = DataLoader(dataset_train,
+                                  config.DATA.BATCH_SIZE,
+                                  sampler=sampler_train,
+                                  num_workers=config.DATA.NUM_WORKERS,
+                                  pin_memory=True,
+                                  drop_last=True)
+
+    dataloader_val = DataLoader(dataset_val,
+                                config.DATA.BATCH_SIZE,
+                                sampler=sampler_val,
+                                num_workers=config.DATA.NUM_WORKERS,
+                                pin_memory=True,
+                                drop_last=False)
+
+    return dataloader_train, dataloader_val

pytorch-caney/pytorch_caney/data/datamodules/mim_datamodule.py

@@ -0,0 +1,80 @@
+from ..datasets.simmim_modis_dataset import MODISDataset
+
+from ..transforms import SimmimTransform
+
+import torch.distributed as dist
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.data._utils.collate import default_collate
+
+
+DATASETS = {
+    'MODIS': MODISDataset,
+}
+
+
+def collate_fn(batch):
+    if not isinstance(batch[0][0], tuple):
+        return default_collate(batch)
+    else:
+        batch_num = len(batch)
+        ret = []
+        for item_idx in range(len(batch[0][0])):
+            if batch[0][0][item_idx] is None:
+                ret.append(None)
+            else:
+                ret.append(default_collate(
+                    [batch[i][0][item_idx] for i in range(batch_num)]))
+        ret.append(default_collate([batch[i][1] for i in range(batch_num)]))
+        return ret
+
+
+def get_dataset_from_dict(dataset_name):
+
+    try:
+
+        dataset_to_use = DATASETS[dataset_name]
+
+    except KeyError:
+
+        error_msg = f"{dataset_name} is not an existing dataset"
+
+        error_msg = f"{error_msg}. Available datasets: {DATASETS.keys()}"
+
+        raise KeyError(error_msg)
+
+    return dataset_to_use
+
+
+def build_mim_dataloader(config, logger):
+
+    transform = SimmimTransform(config)
+
+    logger.info(f'Pre-train data transform:\n{transform}')
+
+    dataset_name = config.DATA.DATASET
+
+    dataset_to_use = get_dataset_from_dict(dataset_name)
+
+    dataset = dataset_to_use(config,
+                             config.DATA.DATA_PATHS,
+                             split="train",
+                             img_size=config.DATA.IMG_SIZE,
+                             transform=transform)
+
+    logger.info(f'Build dataset: train images = {len(dataset)}')
+
+    sampler = DistributedSampler(
+        dataset,
+        num_replicas=dist.get_world_size(),
+        rank=dist.get_rank(),
+        shuffle=True)
+
+    dataloader = DataLoader(dataset,
+                            config.DATA.BATCH_SIZE,
+                            sampler=sampler,
+                            num_workers=config.DATA.NUM_WORKERS,
+                            pin_memory=True,
+                            drop_last=True,
+                            collate_fn=collate_fn)
+
+    return dataloader

pytorch-caney/pytorch_caney/data/datamodules/segmentation_datamodule.py

@@ -0,0 +1,164 @@
+import os
+import logging
+from typing import Any, Union, Optional
+
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import random_split
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning.utilities.cli import DATAMODULE_REGISTRY
+
+from terragpu.ai.deep_learning.datasets.segmentation_dataset \
+    import SegmentationDataset
+
+
+@DATAMODULE_REGISTRY
+class SegmentationDataModule(LightningDataModule):
+
+    def __init__(
+        self,
+
+        # Dataset parameters
+        dataset_dir: str = 'dataset/',
+        images_regex: str = 'dataset/images/*.tif',
+        labels_regex: str = 'dataset/labels/*.tif',
+        generate_dataset: bool = True,
+        tile_size: int = 256,
+        max_patches: Union[float, int] = 100,
+        augment: bool = True,
+        chunks: dict = {'band': 1, 'x': 2048, 'y': 2048},
+        input_bands: list = ['CB', 'B', 'G', 'Y', 'R', 'RE', 'N1', 'N2'],
+        output_bands: list = ['B', 'G', 'R'],
+        seed: int = 24,
+        normalize: bool = True,
+        pytorch: bool = True,
+
+        # Datamodule parameters
+        val_split: float = 0.2,
+        test_split: float = 0.1,
+        num_workers: int = os.cpu_count(),
+        batch_size: int = 32,
+        shuffle: bool = True,
+        pin_memory: bool = False,
+        drop_last: bool = False,
+
+        # Inference parameters
+        raster_regex: str = 'rasters/*.tif',
+
+        *args: Any,
+        **kwargs: Any,
+
+    ) -> None:
+
+        super().__init__(*args, **kwargs)
+
+        # Dataset parameters
+        self.images_regex = images_regex
+        self.labels_regex = labels_regex
+        self.dataset_dir = dataset_dir
+        self.generate_dataset = generate_dataset
+        self.tile_size = tile_size
+        self.max_patches = max_patches
+        self.augment = augment
+        self.chunks = chunks
+        self.input_bands = input_bands
+        self.output_bands = output_bands
+        self.seed = seed
+        self.normalize = normalize
+        self.pytorch = pytorch
+
+        self.val_split = val_split
+        self.test_split = test_split
+        self.raster_regex = raster_regex
+
+        # Performance parameters
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.shuffle = shuffle
+        self.pin_memory = pin_memory
+        self.drop_last = drop_last
+
+    def prepare_data(self):
+        if self.generate_dataset:
+            SegmentationDataset(
+                images_regex=self.images_regex,
+                labels_regex=self.labels_regex,
+                dataset_dir=self.dataset_dir,
+                generate_dataset=self.generate_dataset,
+                tile_size=self.tile_size,
+                max_patches=self.max_patches,
+                augment=self.augment,
+                chunks=self.chunks,
+                input_bands=self.input_bands,
+                output_bands=self.output_bands,
+                seed=self.seed,
+                normalize=self.normalize,
+                pytorch=self.pytorch,
+            )
+
+    def setup(self, stage: Optional[str] = None):
+
+        # Split into train, val, test
+        segmentation_dataset = SegmentationDataset(
+            images_regex=self.images_regex,
+            labels_regex=self.labels_regex,
+            dataset_dir=self.dataset_dir,
+            generate_dataset=False,
+            tile_size=self.tile_size,
+            max_patches=self.max_patches,
+            augment=self.augment,
+            chunks=self.chunks,
+            input_bands=self.input_bands,
+            output_bands=self.output_bands,
+            seed=self.seed,
+            normalize=self.normalize,
+            pytorch=self.pytorch,
+        )
+
+        # Split datasets into train, val, and test sets
+        val_len = round(self.val_split * len(segmentation_dataset))
+        test_len = round(self.test_split * len(segmentation_dataset))
+        train_len = len(segmentation_dataset) - val_len - test_len
+
+        # Initialize datasets
+        self.train_set, self.val_set, self.test_set = random_split(
+            segmentation_dataset, lengths=[train_len, val_len, test_len],
+            generator=torch.Generator().manual_seed(self.seed)
+        )
+        logging.info("Initialized datasets...")
+
+    def train_dataloader(self) -> DataLoader:
+        loader = DataLoader(
+            self.train_set,
+            batch_size=self.batch_size,
+            shuffle=self.shuffle,
+            num_workers=self.num_workers,
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory,
+        )
+        return loader
+
+    def val_dataloader(self) -> DataLoader:
+        loader = DataLoader(
+            self.val_set,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory,
+        )
+        return loader
+
+    def test_dataloader(self) -> DataLoader:
+        loader = DataLoader(
+            self.test_set,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=self.drop_last,
+            pin_memory=self.pin_memory,
+        )
+        return loader
+
+    def predict_dataloader(self) -> DataLoader:
+        raise NotImplementedError

pytorch-caney/pytorch_caney/data/datamodules/simmim_datamodule.py

@@ -0,0 +1,80 @@
+from ..datasets.simmim_modis_dataset import MODISDataset
+
+from ..transforms import SimmimTransform
+
+import torch.distributed as dist
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.data._utils.collate import default_collate
+
+
+DATASETS = {
+    'MODIS': MODISDataset,
+}
+
+
+def collate_fn(batch):
+    if not isinstance(batch[0][0], tuple):
+        return default_collate(batch)
+    else:
+        batch_num = len(batch)
+        ret = []
+        for item_idx in range(len(batch[0][0])):
+            if batch[0][0][item_idx] is None:
+                ret.append(None)
+            else:
+                ret.append(default_collate(
+                    [batch[i][0][item_idx] for i in range(batch_num)]))
+        ret.append(default_collate([batch[i][1] for i in range(batch_num)]))
+        return ret
+
+
+def get_dataset_from_dict(dataset_name):
+
+    try:
+
+        dataset_to_use = DATASETS[dataset_name]
+
+    except KeyError:
+
+        error_msg = f"{dataset_name} is not an existing dataset"
+
+        error_msg = f"{error_msg}. Available datasets: {DATASETS.keys()}"
+
+        raise KeyError(error_msg)
+
+    return dataset_to_use
+
+
+def build_mim_dataloader(config, logger):
+
+    transform = SimmimTransform(config)
+
+    logger.info(f'Pre-train data transform:\n{transform}')
+
+    dataset_name = config.DATA.DATASET
+
+    dataset_to_use = get_dataset_from_dict(dataset_name)
+
+    dataset = dataset_to_use(config,
+                             config.DATA.DATA_PATHS,
+                             split="train",
+                             img_size=config.DATA.IMG_SIZE,
+                             transform=transform)
+
+    logger.info(f'Build dataset: train images = {len(dataset)}')
+
+    sampler = DistributedSampler(
+        dataset,
+        num_replicas=dist.get_world_size(),
+        rank=dist.get_rank(),
+        shuffle=True)
+
+    dataloader = DataLoader(dataset,
+                            config.DATA.BATCH_SIZE,
+                            sampler=sampler,
+                            num_workers=config.DATA.NUM_WORKERS,
+                            pin_memory=True,
+                            drop_last=True,
+                            collate_fn=collate_fn)
+
+    return dataloader