Models¶
ModelABC¶
classDiagram
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
class ArchesWeather {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() tuple
}
ModelABC <|-- ArchesWeather : herits
class BaseModel {
}
ModelABC <|-- BaseModel : herits
class DeepLabV3 {
+bool onnx_supported
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+check_input_shape() None
+forward() torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+get_classification_head() Sequential
+get_decoder() DeepLabV3Decoder
+get_segmentation_head() Sequential
+initialize() None
+initialize_decoder() None
+initialize_head() None
+predict() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- DeepLabV3 : herits
class DeepLabV3Plus {
}
ModelABC <|-- DeepLabV3Plus : herits
class HalfUNet {
+bool onnx_supported
+tuple supported_num_spatial_dims
+int num_spatial_dims
+bool features_last
+ModelType model_type
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- HalfUNet : herits
class IdentityModel {
+int num_spatial_dims
+bool register
+forward() Tensor
}
ModelABC <|-- IdentityModel : herits
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
ModelABC <|-- BaseGraphModel : herits
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
ModelABC <|-- BaseHiGraphModel : herits
class GraphLAM {
+bool register
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+process_step() Tensor
}
ModelABC <|-- GraphLAM : herits
class HiLAM {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
+make_down_gnns() ModuleList
+make_same_gnns() ModuleList
+make_up_gnns() ModuleList
+mesh_down_step() tuple
+mesh_up_step() tuple
}
ModelABC <|-- HiLAM : herits
class HiLAMParallel {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
}
ModelABC <|-- HiLAMParallel : herits
class PanguWeather {
+bool onnx_supported
+Tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() Tuple
}
ModelABC <|-- PanguWeather : herits
class Segformer {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- Segformer : herits
class SwinUNetR {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- SwinUNetR : herits
class CustomUNet {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() Tuple
}
ModelABC <|-- CustomUNet : herits
class UNet {
+int num_spatial_dims
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- UNet : herits
class UNetRPP {
+bool features_last
+bool register
+forward() torch.Tensor | list[torch.Tensor]
+proj_feat() Tensor
+validate_input_shape() tuple
}
ModelABC <|-- UNetRPP : herits
class ViTClassifier {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
ModelABC <|-- ViTClassifier : herits
class VitEncoder {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
ModelABC <|-- VitEncoder : herits
ArchesWeather¶
classDiagram
class ArchesWeather {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- ArchesWeather : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- ArchesWeather : herits
class Module
<<abstract>> Module
Module <|-- ArchesWeather : herits
ArchesWeather model as described in http://arxiv.org/abs/2405.14527. |
BaseModel¶
classDiagram
class BaseModel {
}
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- BaseModel : herits
class Module
<<abstract>> Module
Module <|-- BaseModel : herits
class ArchesWeather {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() tuple
}
BaseModel <|-- ArchesWeather : herits
class DeepLabV3 {
+bool onnx_supported
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+check_input_shape() None
+forward() torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+get_classification_head() Sequential
+get_decoder() DeepLabV3Decoder
+get_segmentation_head() Sequential
+initialize() None
+initialize_decoder() None
+initialize_head() None
+predict() Tensor
+validate_input_shape() tuple
}
BaseModel <|-- DeepLabV3 : herits
class DeepLabV3Plus {
}
BaseModel <|-- DeepLabV3Plus : herits
class HalfUNet {
+bool onnx_supported
+tuple supported_num_spatial_dims
+int num_spatial_dims
+bool features_last
+ModelType model_type
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
BaseModel <|-- HalfUNet : herits
class IdentityModel {
+int num_spatial_dims
+bool register
+forward() Tensor
}
BaseModel <|-- IdentityModel : herits
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
BaseModel <|-- BaseGraphModel : herits
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
BaseModel <|-- BaseHiGraphModel : herits
class GraphLAM {
+bool register
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+process_step() Tensor
}
BaseModel <|-- GraphLAM : herits
class HiLAM {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
+make_down_gnns() ModuleList
+make_same_gnns() ModuleList
+make_up_gnns() ModuleList
+mesh_down_step() tuple
+mesh_up_step() tuple
}
BaseModel <|-- HiLAM : herits
class HiLAMParallel {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
}
BaseModel <|-- HiLAMParallel : herits
class PanguWeather {
+bool onnx_supported
+Tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() Tuple
}
BaseModel <|-- PanguWeather : herits
class Segformer {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() tuple
}
BaseModel <|-- Segformer : herits
class CustomUNet {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() Tuple
}
BaseModel <|-- CustomUNet : herits
class UNet {
+int num_spatial_dims
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
BaseModel <|-- UNet : herits
class UNetRPP {
+bool features_last
+bool register
+forward() torch.Tensor | list[torch.Tensor]
+proj_feat() Tensor
+validate_input_shape() tuple
}
BaseModel <|-- UNetRPP : herits
class ViTClassifier {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
BaseModel <|-- ViTClassifier : herits
class VitEncoder {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
BaseModel <|-- VitEncoder : herits
DeepLabV3¶
classDiagram
class DeepLabV3 {
+bool onnx_supported
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+check_input_shape() None
+forward() torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+get_classification_head() Sequential
+get_decoder() DeepLabV3Decoder
+get_segmentation_head() Sequential
+initialize() None
+initialize_decoder() None
+initialize_head() None
+predict() Tensor
+validate_input_shape() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- DeepLabV3 : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- DeepLabV3 : herits
class Module
<<abstract>> Module
Module <|-- DeepLabV3 : herits
class DeepLabV3Plus {
}
DeepLabV3 <|-- DeepLabV3Plus : herits
DeepLabV3Plus¶
classDiagram
class DeepLabV3Plus {
}
class DeepLabV3 {
+bool onnx_supported
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+check_input_shape() None
+forward() torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+get_classification_head() Sequential
+get_decoder() DeepLabV3Decoder
+get_segmentation_head() Sequential
+initialize() None
+initialize_decoder() None
+initialize_head() None
+predict() Tensor
+validate_input_shape() tuple
}
DeepLabV3 <|-- DeepLabV3Plus : herits
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- DeepLabV3Plus : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- DeepLabV3Plus : herits
class Module
<<abstract>> Module
Module <|-- DeepLabV3Plus : herits
DeepLabV3+ implementation from "Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation". |
HalfUNet¶
classDiagram
class HalfUNet {
+bool onnx_supported
+tuple supported_num_spatial_dims
+int num_spatial_dims
+bool features_last
+ModelType model_type
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- HalfUNet : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- HalfUNet : herits
class Module
<<abstract>> Module
Module <|-- HalfUNet : herits
IdentityModel¶
classDiagram
class IdentityModel {
+int num_spatial_dims
+bool register
+forward() Tensor
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- IdentityModel : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- IdentityModel : herits
class Module
<<abstract>> Module
Module <|-- IdentityModel : herits
Implementation of an identity model. |
BaseGraphModel¶
classDiagram
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- BaseGraphModel : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- BaseGraphModel : herits
class Module
<<abstract>> Module
Module <|-- BaseGraphModel : herits
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
BaseGraphModel <|-- BaseHiGraphModel : herits
class GraphLAM {
+bool register
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+process_step() Tensor
}
BaseGraphModel <|-- GraphLAM : herits
class HiLAM {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
+make_down_gnns() ModuleList
+make_same_gnns() ModuleList
+make_up_gnns() ModuleList
+mesh_down_step() tuple
+mesh_up_step() tuple
}
BaseGraphModel <|-- HiLAM : herits
class HiLAMParallel {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
}
BaseGraphModel <|-- HiLAMParallel : herits
Base (abstract) class for graph-based models building on the encode-process-decode idea. |
BaseHiGraphModel¶
classDiagram
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
BaseGraphModel <|-- BaseHiGraphModel : herits
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- BaseHiGraphModel : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- BaseHiGraphModel : herits
class Module
<<abstract>> Module
Module <|-- BaseHiGraphModel : herits
class HiLAM {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
+make_down_gnns() ModuleList
+make_same_gnns() ModuleList
+make_up_gnns() ModuleList
+mesh_down_step() tuple
+mesh_up_step() tuple
}
BaseHiGraphModel <|-- HiLAM : herits
class HiLAMParallel {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
}
BaseHiGraphModel <|-- HiLAMParallel : herits
Base class for hierarchical graph models. |
GraphLAM¶
classDiagram
class GraphLAM {
+bool register
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+process_step() Tensor
}
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
BaseGraphModel <|-- GraphLAM : herits
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- GraphLAM : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- GraphLAM : herits
class Module
<<abstract>> Module
Module <|-- GraphLAM : herits
Full graph-based LAM model that can be used with different (non-hierarchical )graphs. |
HiLAM¶
classDiagram
class HiLAM {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
+make_down_gnns() ModuleList
+make_same_gnns() ModuleList
+make_up_gnns() ModuleList
+mesh_down_step() tuple
+mesh_up_step() tuple
}
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
BaseHiGraphModel <|-- HiLAM : herits
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
BaseGraphModel <|-- HiLAM : herits
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- HiLAM : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- HiLAM : herits
class Module
<<abstract>> Module
Module <|-- HiLAM : herits
Hierarchical graph model with message passing that goes sequentially down and up the hierarchy during processing. |
HiLAMParallel¶
classDiagram
class HiLAMParallel {
+bool register
+finalize_graph_model() None
+hi_processor_step() tuple
}
class BaseHiGraphModel {
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+get_num_mesh() tuple
+hi_processor_step() tuple
+process_step() Tensor
}
BaseHiGraphModel <|-- HiLAMParallel : herits
class BaseGraphModel {
+int num_spatial_dims
+bool features_last
+embedd_mesh_nodes() Tensor
+finalize_graph_model() None
+forward() Tensor
+get_num_mesh() tuple
+load_graph() None
+process_step() Tensor
}
BaseGraphModel <|-- HiLAMParallel : herits
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- HiLAMParallel : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- HiLAMParallel : herits
class Module
<<abstract>> Module
Module <|-- HiLAMParallel : herits
Version of HiLAM where all message passing in the hierarchical mesh (up, down, inter-level) is ran in paralell. |
PanguWeather¶
classDiagram
class PanguWeather {
+bool onnx_supported
+Tuple supported_num_spatial_dims
+bool features_last
+bool register
+forward() Tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- PanguWeather : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- PanguWeather : herits
class Module
<<abstract>> Module
Module <|-- PanguWeather : herits
PanguWeather network as described in http://arxiv.org/abs/2211.02556 and https://www.nature.com/articles/s41586-023-06185-3. |
Segformer¶
classDiagram
class Segformer {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- Segformer : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- Segformer : herits
class Module
<<abstract>> Module
Module <|-- Segformer : herits
Segformer architecture with extra upsampling in the decoder to match the input image size. |
SwinUNetR¶
classDiagram
class SwinUNetR {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() tuple
}
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- SwinUNetR : herits
class Module
<<abstract>> Module
Module <|-- SwinUNetR : herits
Wrapper around the SwinUNETR from MONAI. |
CustomUNet¶
classDiagram
class CustomUNet {
+int num_spatial_dims
+bool register
+forward() Tensor
+validate_input_shape() Tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- CustomUNet : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- CustomUNet : herits
class Module
<<abstract>> Module
Module <|-- CustomUNet : herits
CustomUNet is a model that allows the user to define a specific configuration, from pretrained weights or not (from ResNet encoders). |
UNet¶
classDiagram
class UNet {
+int num_spatial_dims
+bool register
#_block() Sequential
+forward() Tensor
+validate_input_shape() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- UNet : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- UNet : herits
class Module
<<abstract>> Module
Module <|-- UNet : herits
Returns a UNet architecture, with uninitialised weights, matching desired numbers of input and output channels. |
UNetRPP¶
classDiagram
class UNetRPP {
+bool features_last
+bool register
+forward() torch.Tensor | list[torch.Tensor]
+proj_feat() Tensor
+validate_input_shape() tuple
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- UNetRPP : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- UNetRPP : herits
class Module
<<abstract>> Module
Module <|-- UNetRPP : herits
UNetR++ based on: "Shaker et al., UNETR++: Delving into Efficient and Accurate 3D Medical Image Segmentation". |
ViTClassifier¶
classDiagram
class ViTClassifier {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- ViTClassifier : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- ViTClassifier : herits
class Module
<<abstract>> Module
Module <|-- ViTClassifier : herits
Vision Transformer (ViT) classifier model outputing class probabilities per input sample. |
VitEncoder¶
classDiagram
class VitEncoder {
+bool onnx_supported
+tuple supported_num_spatial_dims
+bool features_last
+ModelType model_type
+int num_spatial_dims
+bool register
+forward() Tensor
}
class BaseModel {
}
<<abstract>> BaseModel
BaseModel <|-- VitEncoder : herits
class ModelABC {
+bool register
+int in_channels
+int out_channels
+tuple input_shape
+check_required_attributes() None
}
<<abstract>> ModelABC
ModelABC <|-- VitEncoder : herits
class Module
<<abstract>> Module
Module <|-- VitEncoder : herits
ViT vision encoder for multimodal LLMs. |