Models

models

Classes

DeepForestModule

Bases: LightningModule

Source code in tree_detection_framework/detection/models.py

class DeepForestModule(lightning.LightningModule):
    def __init__(
        self,
        use_hugging_face_weights: bool = True,
        param_dict: Optional[Dict[str, Any]] = {},
    ):
        """_summary_

        Args:
            use_hugging_face_weights (bool, optional):
                Should the model and weights be donwloaded from the Hugging Face repository. If
                True, setting param_dict will result in an error. Defaults to True.
            param_dict (Optional[Dict[str, Any]], optional):
                Configuration parameters to provide finer control over the model. Cannot be used
                with use_hugging_face_weights=True. Defaults to {}.

        Raises:
            ValueError: If both use_hugging_face_weights and param_dict are set
        """
        super().__init__()
        self.param_dict = param_dict

        # Determine how to obtain the weights
        if use_hugging_face_weights:
            # Error if the user tried to use both options
            if len(param_dict) > 0:
                raise ValueError(
                    "Setting the `param_dict` and `use_hugging_face_weights`=True are mutually exclusive. Please choose one option."
                )
            # Create the architecture
            model = deepforest_main.deepforest()
            # Load a pretrained tree detection model from Hugging Face
            model.load_model(model_name="weecology/deepforest-tree", revision="main")
            # The definition of the model here contains additional metrics that do not match what we
            # need. Take only the prediction model portion.
            self.model = model.model

        else:
            if param_dict["backbone"] == "retinanet":
                retinanet = RetinaNetModel(param_dict)
            else:
                raise ValueError("Only 'retinanet' backbone is currently supported.")

            self.model = retinanet.create_model()
            self.use_release()

    def use_release(self, check_release=True):
        """Use the latest DeepForest model release from github and load model.
        Optionally download if release doesn't exist.
        Args:
            check_release (logical): whether to check github for a model recent release.
            In cases where you are hitting the github API rate limit, set to False and any local model will be downloaded.
            If no model has been downloaded an error will raise.
        """
        # Download latest model from github release
        release_tag, self.release_state_dict = use_release_df(
            check_release=check_release
        )
        self.model.load_state_dict(torch.load(self.release_state_dict))

        # load saved model and tag release
        self.__release_version__ = release_tag
        print("Loading pre-built model: {}".format(release_tag))

    def forward(
        self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None
    ) -> Dict[str, Tensor]:
        """Calls the model's forward method.
        Args:
            images (list[Tensor]): Images to be processed
            targets (list[Dict[Tensor]]): Ground-truth boxes present in the image (optional)

        Returns:
            result (list[BoxList] or dict[Tensor]):
                The output from the model.
                During training, it returns a dict[Tensor] which contains the losses.
        """
        # Move the data to the same device as the model
        images = images.to(self.device)
        return self.model.forward(images, targets=targets)  # Model specific forward

    def training_step(self, batch):
        # Ensure model is in train mode
        self.model.train()
        device = next(self.model.parameters()).device

        # Image is expected to be a list of tensors, each of shape [C, H, W] in 0-1 range.
        image_batch = (batch["image"][:, :3, :, :] / 255.0).to(device)
        image_batch_list = [image for image in image_batch]

        # To store every image's target - a dictionary containing `boxes` and `labels`
        targets = []
        for tile in batch["bounding_boxes"]:
            # Convert from list to FloatTensor[N, 4]
            boxes_tensor = torch.tensor(tile, dtype=torch.float32).to(device)
            # Need to remove boxes that go out-of-bounds. Has negative values.
            valid_mask = (boxes_tensor >= 0).all(dim=1)
            filtered_boxes_tensor = boxes_tensor[valid_mask]
            # Create a label tensor. Single class for now.
            class_labels = torch.zeros(
                filtered_boxes_tensor.shape[0], dtype=torch.int64
            ).to(device)
            # Dictionary for the tile
            d = {"boxes": filtered_boxes_tensor, "labels": class_labels}
            targets.append(d)

        loss_dict = self.forward(image_batch_list, targets=targets)

        final_loss = sum([loss for loss in loss_dict.values()])
        print("loss: ", final_loss)
        return final_loss

    def configure_optimizers(self):
        # similar to the one in deepforest
        optimizer = optim.SGD(
            self.model.parameters(), lr=self.param_dict["train"]["lr"], momentum=0.9
        )

        # TODO: Setup lr_scheduler
        # TODO: Return 'optimizer', 'lr_scheduler', 'monitor' when validation data is set

        return optimizer

Methods:

__init__(use_hugging_face_weights=True, param_dict={})

summary

Parameters:

Name	Type	Description	Default
use_hugging_face_weights	bool	Should the model and weights be donwloaded from the Hugging Face repository. If True, setting param_dict will result in an error. Defaults to True.	True
param_dict	Optional[Dict[str, Any]]	Configuration parameters to provide finer control over the model. Cannot be used with use_hugging_face_weights=True. Defaults to {}.	{}

Raises:

Type	Description
ValueError	If both use_hugging_face_weights and param_dict are set

Source code in tree_detection_framework/detection/models.py

def __init__(
    self,
    use_hugging_face_weights: bool = True,
    param_dict: Optional[Dict[str, Any]] = {},
):
    """_summary_

    Args:
        use_hugging_face_weights (bool, optional):
            Should the model and weights be donwloaded from the Hugging Face repository. If
            True, setting param_dict will result in an error. Defaults to True.
        param_dict (Optional[Dict[str, Any]], optional):
            Configuration parameters to provide finer control over the model. Cannot be used
            with use_hugging_face_weights=True. Defaults to {}.

    Raises:
        ValueError: If both use_hugging_face_weights and param_dict are set
    """
    super().__init__()
    self.param_dict = param_dict

    # Determine how to obtain the weights
    if use_hugging_face_weights:
        # Error if the user tried to use both options
        if len(param_dict) > 0:
            raise ValueError(
                "Setting the `param_dict` and `use_hugging_face_weights`=True are mutually exclusive. Please choose one option."
            )
        # Create the architecture
        model = deepforest_main.deepforest()
        # Load a pretrained tree detection model from Hugging Face
        model.load_model(model_name="weecology/deepforest-tree", revision="main")
        # The definition of the model here contains additional metrics that do not match what we
        # need. Take only the prediction model portion.
        self.model = model.model

    else:
        if param_dict["backbone"] == "retinanet":
            retinanet = RetinaNetModel(param_dict)
        else:
            raise ValueError("Only 'retinanet' backbone is currently supported.")

        self.model = retinanet.create_model()
        self.use_release()

forward(images, targets=None)

Calls the model's forward method. Args: images (list[Tensor]): Images to be processed targets (list[Dict[Tensor]]): Ground-truth boxes present in the image (optional)

Returns:

Name	Type	Description
result	list[BoxList] or dict[Tensor]	The output from the model. During training, it returns a dict[Tensor] which contains the losses.

Source code in tree_detection_framework/detection/models.py

def forward(
    self, images: List[Tensor], targets: Optional[List[Dict[str, Tensor]]] = None
) -> Dict[str, Tensor]:
    """Calls the model's forward method.
    Args:
        images (list[Tensor]): Images to be processed
        targets (list[Dict[Tensor]]): Ground-truth boxes present in the image (optional)

    Returns:
        result (list[BoxList] or dict[Tensor]):
            The output from the model.
            During training, it returns a dict[Tensor] which contains the losses.
    """
    # Move the data to the same device as the model
    images = images.to(self.device)
    return self.model.forward(images, targets=targets)  # Model specific forward

use_release(check_release=True)

Use the latest DeepForest model release from github and load model. Optionally download if release doesn't exist. Args: check_release (logical): whether to check github for a model recent release. In cases where you are hitting the github API rate limit, set to False and any local model will be downloaded. If no model has been downloaded an error will raise.

Source code in tree_detection_framework/detection/models.py

def use_release(self, check_release=True):
    """Use the latest DeepForest model release from github and load model.
    Optionally download if release doesn't exist.
    Args:
        check_release (logical): whether to check github for a model recent release.
        In cases where you are hitting the github API rate limit, set to False and any local model will be downloaded.
        If no model has been downloaded an error will raise.
    """
    # Download latest model from github release
    release_tag, self.release_state_dict = use_release_df(
        check_release=check_release
    )
    self.model.load_state_dict(torch.load(self.release_state_dict))

    # load saved model and tag release
    self.__release_version__ = release_tag
    print("Loading pre-built model: {}".format(release_tag))

Detectree2Module

Source code in tree_detection_framework/detection/models.py

class Detectree2Module:
    def __init__(self, param_dict: Optional[Dict[str, Any]] = None):
        if DETECTRON2_AVAILABLE is False:
            raise ImportError(
                "detectron2 is not installed. Please install it to use this module."
            )
        super().__init__()
        # If param_dict is not provided, ensure it is an empty dictionary
        self.param_dict = param_dict or {}
        self.cfg = self.setup_cfg(**self.param_dict)

    def setup_cfg(
        self,
        base_model: str = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml",
        trains=("trees_train",),
        tests=("trees_val",),
        update_model=None,
        workers=2,
        ims_per_batch=2,
        gamma=0.1,
        backbone_freeze=3,
        warm_iter=120,
        momentum=0.9,
        batch_size_per_im=1024,
        base_lr=0.0003389,
        weight_decay=0.001,
        max_iter=1000,
        num_classes=1,
        eval_period=100,
        out_dir="./train_outputs",
        resize=True,
    ):
        """Set up config object.
        Args:
            base_model: base pre-trained model from detectron2 model_zoo
            trains: names of registered data to use for training
            tests: names of registered data to use for evaluating models
            update_model: updated pre-trained model from detectree2 model_garden
            workers: number of workers for dataloader
            ims_per_batch: number of images per batch
            gamma: gamma for learning rate scheduler
            backbone_freeze: backbone layer to freeze
            warm_iter: number of iterations for warmup
            momentum: momentum for optimizer
            batch_size_per_im: batch size per image
            base_lr: base learning rate
            weight_decay: weight decay for optimizer
            max_iter: maximum number of iterations
            num_classes: number of classes
            eval_period: number of iterations between evaluations
            out_dir: directory to save outputs
            resize: whether to resize input images
        """
        # Initialize configuration
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file(base_model))

        # Assign values, prioritizing those in param_dict
        cfg.DATASETS.TRAIN = self.param_dict.get("trains", trains)
        cfg.DATASETS.TEST = self.param_dict.get("tests", tests)
        cfg.DATALOADER.NUM_WORKERS = self.param_dict.get("workers", workers)
        cfg.SOLVER.IMS_PER_BATCH = self.param_dict.get("ims_per_batch", ims_per_batch)
        cfg.SOLVER.GAMMA = self.param_dict.get("gamma", gamma)
        cfg.MODEL.BACKBONE.FREEZE_AT = self.param_dict.get(
            "backbone_freeze", backbone_freeze
        )
        cfg.SOLVER.WARMUP_ITERS = self.param_dict.get("warm_iter", warm_iter)
        cfg.SOLVER.MOMENTUM = self.param_dict.get("momentum", momentum)
        cfg.MODEL.RPN.BATCH_SIZE_PER_IMAGE = self.param_dict.get(
            "batch_size_per_im", batch_size_per_im
        )
        cfg.SOLVER.WEIGHT_DECAY = self.param_dict.get("weight_decay", weight_decay)
        cfg.SOLVER.BASE_LR = self.param_dict.get("base_lr", base_lr)
        cfg.OUTPUT_DIR = self.param_dict.get("out_dir", out_dir)
        cfg.SOLVER.MAX_ITER = self.param_dict.get("max_iter", max_iter)
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = self.param_dict.get(
            "num_classes", num_classes
        )
        cfg.TEST.EVAL_PERIOD = self.param_dict.get("eval_period", eval_period)
        cfg.RESIZE = self.param_dict.get("resize", resize)
        cfg.INPUT.MIN_SIZE_TRAIN = 1000

        # Create output directory if it doesn't exist
        Path(cfg.OUTPUT_DIR).mkdir(parents=True, exist_ok=True)

        # Set model weights
        cfg.MODEL.WEIGHTS = self.param_dict.get(
            "update_model", update_model
        ) or model_zoo.get_checkpoint_url(base_model)

        return cfg

Methods:

setup_cfg(base_model='COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml', trains=('trees_train',), tests=('trees_val',), update_model=None, workers=2, ims_per_batch=2, gamma=0.1, backbone_freeze=3, warm_iter=120, momentum=0.9, batch_size_per_im=1024, base_lr=0.0003389, weight_decay=0.001, max_iter=1000, num_classes=1, eval_period=100, out_dir='./train_outputs', resize=True)

Set up config object. Args: base_model: base pre-trained model from detectron2 model_zoo trains: names of registered data to use for training tests: names of registered data to use for evaluating models update_model: updated pre-trained model from detectree2 model_garden workers: number of workers for dataloader ims_per_batch: number of images per batch gamma: gamma for learning rate scheduler backbone_freeze: backbone layer to freeze warm_iter: number of iterations for warmup momentum: momentum for optimizer batch_size_per_im: batch size per image base_lr: base learning rate weight_decay: weight decay for optimizer max_iter: maximum number of iterations num_classes: number of classes eval_period: number of iterations between evaluations out_dir: directory to save outputs resize: whether to resize input images

Source code in tree_detection_framework/detection/models.py

def setup_cfg(
    self,
    base_model: str = "COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml",
    trains=("trees_train",),
    tests=("trees_val",),
    update_model=None,
    workers=2,
    ims_per_batch=2,
    gamma=0.1,
    backbone_freeze=3,
    warm_iter=120,
    momentum=0.9,
    batch_size_per_im=1024,
    base_lr=0.0003389,
    weight_decay=0.001,
    max_iter=1000,
    num_classes=1,
    eval_period=100,
    out_dir="./train_outputs",
    resize=True,
):
    """Set up config object.
    Args:
        base_model: base pre-trained model from detectron2 model_zoo
        trains: names of registered data to use for training
        tests: names of registered data to use for evaluating models
        update_model: updated pre-trained model from detectree2 model_garden
        workers: number of workers for dataloader
        ims_per_batch: number of images per batch
        gamma: gamma for learning rate scheduler
        backbone_freeze: backbone layer to freeze
        warm_iter: number of iterations for warmup
        momentum: momentum for optimizer
        batch_size_per_im: batch size per image
        base_lr: base learning rate
        weight_decay: weight decay for optimizer
        max_iter: maximum number of iterations
        num_classes: number of classes
        eval_period: number of iterations between evaluations
        out_dir: directory to save outputs
        resize: whether to resize input images
    """
    # Initialize configuration
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(base_model))

    # Assign values, prioritizing those in param_dict
    cfg.DATASETS.TRAIN = self.param_dict.get("trains", trains)
    cfg.DATASETS.TEST = self.param_dict.get("tests", tests)
    cfg.DATALOADER.NUM_WORKERS = self.param_dict.get("workers", workers)
    cfg.SOLVER.IMS_PER_BATCH = self.param_dict.get("ims_per_batch", ims_per_batch)
    cfg.SOLVER.GAMMA = self.param_dict.get("gamma", gamma)
    cfg.MODEL.BACKBONE.FREEZE_AT = self.param_dict.get(
        "backbone_freeze", backbone_freeze
    )
    cfg.SOLVER.WARMUP_ITERS = self.param_dict.get("warm_iter", warm_iter)
    cfg.SOLVER.MOMENTUM = self.param_dict.get("momentum", momentum)
    cfg.MODEL.RPN.BATCH_SIZE_PER_IMAGE = self.param_dict.get(
        "batch_size_per_im", batch_size_per_im
    )
    cfg.SOLVER.WEIGHT_DECAY = self.param_dict.get("weight_decay", weight_decay)
    cfg.SOLVER.BASE_LR = self.param_dict.get("base_lr", base_lr)
    cfg.OUTPUT_DIR = self.param_dict.get("out_dir", out_dir)
    cfg.SOLVER.MAX_ITER = self.param_dict.get("max_iter", max_iter)
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = self.param_dict.get(
        "num_classes", num_classes
    )
    cfg.TEST.EVAL_PERIOD = self.param_dict.get("eval_period", eval_period)
    cfg.RESIZE = self.param_dict.get("resize", resize)
    cfg.INPUT.MIN_SIZE_TRAIN = 1000

    # Create output directory if it doesn't exist
    Path(cfg.OUTPUT_DIR).mkdir(parents=True, exist_ok=True)

    # Set model weights
    cfg.MODEL.WEIGHTS = self.param_dict.get(
        "update_model", update_model
    ) or model_zoo.get_checkpoint_url(base_model)

    return cfg

RetinaNetModel

A backbone class for DeepForest

Source code in tree_detection_framework/detection/models.py

class RetinaNetModel:
    """A backbone class for DeepForest"""

    def __init__(self, param_dict):
        self.param_dict = param_dict

    def load_backbone(self):
        """A torch vision retinanet model"""
        backbone = torchvision.models.detection.retinanet_resnet50_fpn(
            weights=RetinaNet_ResNet50_FPN_Weights.COCO_V1
        )

        return backbone

    def create_anchor_generator(
        self, sizes=((8, 16, 32, 64, 128, 256, 400),), aspect_ratios=((0.5, 1.0, 2.0),)
    ):
        """Create anchor box generator as a function of sizes and aspect ratios"""
        anchor_generator = AnchorGenerator(sizes=sizes, aspect_ratios=aspect_ratios)

        return anchor_generator

    def create_model(self):
        """Create a retinanet model

        Returns:
            model: a pytorch nn module
        """
        resnet = self.load_backbone()
        backbone = resnet.backbone

        model = RetinaNet(backbone=backbone, num_classes=self.param_dict["num_classes"])
        # TODO: do we want to set model.nms_thresh and model.score_thresh?

        return model

Methods:

create_anchor_generator(sizes=((8, 16, 32, 64, 128, 256, 400),), aspect_ratios=((0.5, 1.0, 2.0),))

Create anchor box generator as a function of sizes and aspect ratios

Source code in tree_detection_framework/detection/models.py

def create_anchor_generator(
    self, sizes=((8, 16, 32, 64, 128, 256, 400),), aspect_ratios=((0.5, 1.0, 2.0),)
):
    """Create anchor box generator as a function of sizes and aspect ratios"""
    anchor_generator = AnchorGenerator(sizes=sizes, aspect_ratios=aspect_ratios)

    return anchor_generator

create_model()

Create a retinanet model

Returns:

Name	Type	Description
model		a pytorch nn module

Source code in tree_detection_framework/detection/models.py

def create_model(self):
    """Create a retinanet model

    Returns:
        model: a pytorch nn module
    """
    resnet = self.load_backbone()
    backbone = resnet.backbone

    model = RetinaNet(backbone=backbone, num_classes=self.param_dict["num_classes"])
    # TODO: do we want to set model.nms_thresh and model.score_thresh?

    return model

load_backbone()

A torch vision retinanet model

Source code in tree_detection_framework/detection/models.py

def load_backbone(self):
    """A torch vision retinanet model"""
    backbone = torchvision.models.detection.retinanet_resnet50_fpn(
        weights=RetinaNet_ResNet50_FPN_Weights.COCO_V1
    )

    return backbone

TCDModule

Configuration module for the TCD (Restor Foundation) Mask R-CNN model. Source: https://github.com/Restor-Foundation/tcd

Builds a Detectron2 config matching the settings used to train the TCD Mask R-CNN R50-FPN checkpoint (restor/tcd-mask-rcnn-r50 on HuggingFace Hub). The weights path is stored in cfg.MODEL.WEIGHTS; TCDDetector resolves HuggingFace repo IDs to local paths at load time.

Source code in tree_detection_framework/detection/models.py

class TCDModule:
    """Configuration module for the TCD (Restor Foundation) Mask R-CNN model.
    Source: https://github.com/Restor-Foundation/tcd

    Builds a Detectron2 config matching the settings used to train the TCD Mask R-CNN
    R50-FPN checkpoint (`restor/tcd-mask-rcnn-r50` on HuggingFace Hub). The weights
    path is stored in cfg.MODEL.WEIGHTS; TCDDetector resolves HuggingFace repo IDs to
    local paths at load time.
    """

    _ARCHITECTURE = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
    _HF_CHECKPOINT = "restor/tcd-mask-rcnn-r50"

    def __init__(
        self,
        weights: Optional[str] = None,
        score_thresh_test: float = 0.2,
        max_detections: int = 512,
        min_size_test: int = 0,
        max_size_test: int = 2048,
        num_classes: int = 2,
        device: str = "cuda",
    ):
        """
        Args:
            weights: Local path to checkpoint or HuggingFace repo ID. Defaults to
                `restor/tcd-mask-rcnn-r50`.
            score_thresh_test: Minimum confidence threshold for detections.
            max_detections: Maximum number of detections per image.
            min_size_test: Minimum input size for inference (0 = unconstrained).
            max_size_test: Maximum input size for inference.
            num_classes: Number of output classes (2 for TCD: canopy and tree).
            device: Torch device string ("cuda" or "cpu").
        """
        if DETECTRON2_AVAILABLE is False:
            raise ImportError(
                "detectron2 is not installed. Please install it to use TCDMaskRCNNModule."
            )
        self.cfg = self._setup_cfg(
            weights=weights or self._HF_CHECKPOINT,
            score_thresh_test=score_thresh_test,
            max_detections=max_detections,
            min_size_test=min_size_test,
            max_size_test=max_size_test,
            num_classes=num_classes,
            device=device,
        )

    def _setup_cfg(
        self,
        weights: str,
        score_thresh_test: float,
        max_detections: int,
        min_size_test: int,
        max_size_test: int,
        num_classes: int,
        device: str,
    ):
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file(self._ARCHITECTURE))

        # TCD inference settings (from tcd/config/model/detectron2/detectron_mask_rcnn.yaml)
        cfg.INPUT.MIN_SIZE_TEST = min_size_test
        cfg.INPUT.MAX_SIZE_TEST = max_size_test
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thresh_test
        cfg.MODEL.RPN.PRE_NMS_TOPK_TEST = 512
        cfg.MODEL.RPN.POST_NMS_TOPK_TEST = 512
        cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[32], [64], [128], [256], [512]]
        cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS = [[0.5, 1.0, 2.0]]
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = num_classes
        cfg.MODEL.DEVICE = device
        cfg.TEST.DETECTIONS_PER_IMAGE = max_detections
        cfg.TEST.AUG.ENABLED = False

        cfg.DATASETS.TRAIN = ()
        cfg.DATASETS.TEST = ()
        cfg.MODEL.WEIGHTS = weights

        return cfg

Methods:

__init__(weights=None, score_thresh_test=0.2, max_detections=512, min_size_test=0, max_size_test=2048, num_classes=2, device='cuda')

Parameters:

Name	Type	Description	Default
weights	Optional[str]	Local path to checkpoint or HuggingFace repo ID. Defaults to restor/tcd-mask-rcnn-r50.	None
score_thresh_test	float	Minimum confidence threshold for detections.	0.2
max_detections	int	Maximum number of detections per image.	512
min_size_test	int	Minimum input size for inference (0 = unconstrained).	0
max_size_test	int	Maximum input size for inference.	2048
num_classes	int	Number of output classes (2 for TCD: canopy and tree).	2
device	str	Torch device string ("cuda" or "cpu").	'cuda'

Source code in tree_detection_framework/detection/models.py

def __init__(
    self,
    weights: Optional[str] = None,
    score_thresh_test: float = 0.2,
    max_detections: int = 512,
    min_size_test: int = 0,
    max_size_test: int = 2048,
    num_classes: int = 2,
    device: str = "cuda",
):
    """
    Args:
        weights: Local path to checkpoint or HuggingFace repo ID. Defaults to
            `restor/tcd-mask-rcnn-r50`.
        score_thresh_test: Minimum confidence threshold for detections.
        max_detections: Maximum number of detections per image.
        min_size_test: Minimum input size for inference (0 = unconstrained).
        max_size_test: Maximum input size for inference.
        num_classes: Number of output classes (2 for TCD: canopy and tree).
        device: Torch device string ("cuda" or "cpu").
    """
    if DETECTRON2_AVAILABLE is False:
        raise ImportError(
            "detectron2 is not installed. Please install it to use TCDMaskRCNNModule."
        )
    self.cfg = self._setup_cfg(
        weights=weights or self._HF_CHECKPOINT,
        score_thresh_test=score_thresh_test,
        max_detections=max_detections,
        min_size_test=min_size_test,
        max_size_test=max_size_test,
        num_classes=num_classes,
        device=device,
    )

Functions: