implement SAM backbone with huggingface transformers package (#301)

* non-context vit test passing

* all vit tests passing

* all tests passing

* rebuild docs

Author: themattinthehatt

docs/api/lightning_pose.models.backbones.vit_img_encoder.resample_abs_pos_embed_nhwc.rst

@@ -1,16 +1,16 @@
-ImageEncoderViT_FT
+SamVisionEncoderHF
 ==================
 
-.. currentmodule:: lightning_pose.models.backbones.vit_img_encoder
+.. currentmodule:: lightning_pose.models.backbones.vit_sam
 
-.. autoclass:: ImageEncoderViT_FT
+.. autoclass:: SamVisionEncoderHF
    :show-inheritance:
 
    .. rubric:: Methods Summary
 
    .. autosummary::
 
-      ~ImageEncoderViT_FT.forward
+      ~SamVisionEncoderHF.forward
 
    .. rubric:: Methods Documentation
 

…s.vit_img_encoder.ImageEncoderViT_FT.rst …backbones.vit_sam.SamVisionEncoderHF.rstdocs/api/lightning_pose.models.backbones.vit_img_encoder.ImageEncoderViT_FT.rst renamed to docs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoderHF.rst docs/api/lightning_pose.models.backbones.vit_img_encoder.ImageEncoderViT_FT.rst renamed to docs/api/lightning_pose.models.backbones.vit_sam.SamVisionEncoderHF.rst

@@ -0,0 +1,6 @@
+load_sam_vision_encoder_hf
+==========================
+
+.. currentmodule:: lightning_pose.models.backbones.vit_sam
+
+.. autofunction:: load_sam_vision_encoder_hf

docs/api/lightning_pose.models.backbones.vit_sam.load_sam_vision_encoder_hf.rst

@@ -4,7 +4,7 @@ lightning\_pose.models.backbones
 .. automodapi:: lightning_pose.models.backbones.torchvision
    :no-inheritance-diagram:
 
-.. automodapi:: lightning_pose.models.backbones.vit_img_encoder
+.. automodapi:: lightning_pose.models.backbones.vit_sam
    :no-inheritance-diagram:
 
 .. automodapi:: lightning_pose.models.backbones.vits

docs/modules/lightning_pose.models.backbones.rst

@@ -25,7 +25,7 @@ streamlit
 tensorboard
 torchtyping
 torchvision
+transformers
 typeguard
 typing
 nvidia-dali-cuda110
-segment_anything @ git+https://github.com/facebookresearch/segment-anything.git

docs/requirements.txt

@@ -207,7 +207,7 @@ The following parameters relate to model architecture and unsupervised losses.
     * efficientnet_b0: EfficientNet-B0 pretrained on ImageNet
     * efficientnet_b1: EfficientNet-B1 pretrained on ImageNet
     * efficientnet_b2: EfficientNet-B2 pretrained on ImageNet
-    * vit_b_sam: Segment Anything Model (Vision Transformer Base)
+    * vitb_sam: Segment Anything Model (Vision Transformer Base)
 
   Note: the file size for a single ResNet-50 network is approximately 275 MB.
 

docs/source/user_guide/config_file.rst

@@ -1,4 +1,6 @@
+import importlib.metadata
 from pathlib import Path
+from typing import Any
 
 from omegaconf import OmegaConf
 
@@ -8,9 +10,6 @@
 
 # Hacky way to get version from pypackage.toml.
 # Adapted from: https://github.com/python-poetry/poetry/issues/273#issuecomment-1877789967
-from typing import Any
-import importlib.metadata
-from pathlib import Path
 
 __package_version = "unknown"
 
@@ -33,9 +32,10 @@ def __get_package_version() -> str:
         # Fall back on getting it from a local pyproject.toml.
         # This works in a development environment where the
         # package has not been installed from a distribution.
-        import toml
         import warnings
 
+        import toml
+
         warnings.warn(
             "lightning-pose not pip-installed, getting version from pyproject.toml."
         )

lightning_pose/__init__.py

@@ -7,10 +7,7 @@
 from omegaconf import DictConfig, OmegaConf
 
 from lightning_pose.api.model_config import ModelConfig
-from lightning_pose.data.datatypes import (
-    MultiviewPredictionResult,
-    PredictionResult,
-)
+from lightning_pose.data.datatypes import MultiviewPredictionResult, PredictionResult
 from lightning_pose.models import ALLOWED_MODELS
 from lightning_pose.utils import io as io_utils
 from lightning_pose.utils.predictions import (

lightning_pose/api/model.py

@@ -17,7 +17,11 @@
 import seaborn as sns
 import streamlit as st
 
-from lightning_pose.apps.plots import get_y_label, make_plotly_catplot, make_plotly_scatterplot
+from lightning_pose.apps.plots import (
+    get_y_label,
+    make_plotly_catplot,
+    make_plotly_scatterplot,
+)
 from lightning_pose.apps.utils import (
     build_precomputed_metrics_df,
     get_df_box,

lightning_pose/apps/labeled_frame_diagnostics.py

@@ -16,7 +16,11 @@
 import pandas as pd
 import streamlit as st
 
-from lightning_pose.apps.plots import get_y_label, make_plotly_catplot, plot_precomputed_traces
+from lightning_pose.apps.plots import (
+    get_y_label,
+    make_plotly_catplot,
+    plot_precomputed_traces,
+)
 from lightning_pose.apps.utils import (
     build_precomputed_metrics_df,
     concat_dfs,