Predictors API

culicidaelab.predictors

__all__ = ['MosquitoClassifier', 'MosquitoDetector', 'MosquitoSegmenter', 'ModelWeightsManager'] module-attribute

MosquitoClassifier

Classifies mosquito species from an image using a FastAI model.

This class provides methods to load a pre-trained model, predict species from single or batches of images, evaluate model performance, and visualize the classification results.

Parameters:

Name	Type	Description	Default
settings	Settings	The main settings object for the library, which contains configuration for paths, models, and species.	required
load_model	bool	If True, the model weights are loaded immediately upon initialization. Defaults to False.	False

Attributes:

Name	Type	Description
arch	str	The model architecture (e.g., 'convnext_tiny').
data_dir	Path	The directory where datasets are stored.
species_map	dict[int, str]	A mapping from class indices to species names.
num_classes	int	The total number of species classes.
learner	int	The loaded FastAI learner object, available after load_model().

Source code in culicidaelab/predictors/classifier.py

class MosquitoClassifier(
    BasePredictor[ClassificationPredictionType, ClassificationGroundTruthType],
):
    """Classifies mosquito species from an image using a FastAI model.

    This class provides methods to load a pre-trained model, predict species
    from single or batches of images, evaluate model performance, and visualize
    the classification results.

    Args:
        settings (Settings): The main settings object for the library, which
            contains configuration for paths, models, and species.
        load_model (bool, optional): If True, the model weights are loaded
            immediately upon initialization. Defaults to False.

    Attributes:
        arch (str): The model architecture (e.g., 'convnext_tiny').
        data_dir (Path): The directory where datasets are stored.
        species_map (dict[int, str]): A mapping from class indices to species names.
        num_classes (int): The total number of species classes.
        learner: The loaded FastAI learner object, available after `load_model()`.
    """

    def __init__(self, settings: Settings, load_model: bool = False) -> None:
        """Initializes the MosquitoClassifier."""
        provider_service = ProviderService(settings)
        weights_manager = ModelWeightsManager(
            settings=settings,
            provider_service=provider_service,
        )
        super().__init__(
            settings=settings,
            predictor_type="classifier",
            weights_manager=weights_manager,
            load_model=load_model,
        )

        self.arch: str | None = self.config.model_arch
        self.data_dir: Path = self.settings.dataset_dir
        self.species_map: dict[int, str] = self.settings.species_config.species_map
        self.num_classes: int = len(self.species_map)

    def get_class_index(self, species_name: str) -> int | None:
        """Retrieves the class index for a given species name.

        Args:
            species_name (str): The name of the species.

        Returns:
            int | None: The corresponding class index if found, otherwise None.
        """
        return self.settings.species_config.get_index_by_species(species_name)

    def get_species_names(self) -> list[str]:
        """Gets a sorted list of all species names known to the classifier.

        The list is ordered by the class index.

        Returns:
            list[str]: A list of species names.
        """
        return [self.species_map[i] for i in sorted(self.species_map.keys())]

    def predict(
        self,
        input_data: np.ndarray,
        **kwargs: Any,
    ) -> ClassificationPredictionType:
        """Classifies the mosquito species in a single image.

        Args:
            input_data (np.ndarray): An input image as a NumPy array with a
                shape of (H, W, 3) in RGB format. Values can be uint8
                [0, 255] or float [0, 1].
            **kwargs (Any): Additional arguments (not used).

        Returns:
            ClassificationPredictionType: A list of (species_name, confidence)
            tuples, sorted in descending order of confidence.

        Raises:
            RuntimeError: If the model has not been loaded.
            ValueError: If the input data has an invalid shape, data type,
                or is not a valid image.
        """
        if not self.model_loaded:
            raise RuntimeError(
                "Model is not loaded. Call load_model() or use a context manager.",
            )

        if not isinstance(input_data, np.ndarray):
            input_data = np.array(input_data)

        if input_data.ndim != 3 or input_data.shape[2] != 3:
            raise ValueError(f"Expected 3D RGB image, got shape: {input_data.shape}")

        if input_data.dtype == np.uint8:
            image = Image.fromarray(input_data)
        elif input_data.dtype in [np.float32, np.float64]:
            image = Image.fromarray((input_data * 255).astype(np.uint8))
        else:
            raise ValueError(f"Unsupported dtype: {input_data.dtype}")

        with set_posix_windows():
            _, _, probabilities = self.learner.predict(image)

        species_probs = []
        for idx, prob in enumerate(probabilities):
            species_name = self.species_map.get(idx, f"unknown_{idx}")
            species_probs.append((species_name, float(prob)))

        species_probs.sort(key=lambda x: x[1], reverse=True)
        return species_probs

    def predict_batch(
        self,
        input_data_batch: list[Any],
        show_progress: bool = False,
        **kwargs: Any,
    ) -> list[ClassificationPredictionType]:
        """Classifies mosquito species in a batch of images.

        Note: This method currently iterates and calls `predict` for each image.
        True batch processing is not yet implemented.

        Args:
            input_data_batch (list[Any]): A list of input images, where each
                image is a NumPy array.
            show_progress (bool, optional): If True, a progress bar is displayed.
                Defaults to False.
            **kwargs (Any): Additional arguments passed to `predict`.

        Returns:
            list[ClassificationPredictionType]: A list of prediction results,
            where each result corresponds to an input image.
        """
        results = []
        for img in input_data_batch:
            results.append(self.predict(img, **kwargs))
        return results

    def visualize(
        self,
        input_data: np.ndarray,
        predictions: ClassificationPredictionType,
        save_path: str | Path | None = None,
    ) -> np.ndarray:
        """Overlays classification results on an image.

        This method draws the top-k predictions and their confidence scores
        onto the input image.

        Args:
            input_data (np.ndarray): The original image (H, W, 3) as a NumPy array.
            predictions (ClassificationPredictionType): The prediction output from
                the `predict` method.
            save_path (str | Path | None, optional): If provided, the visualized
                image will be saved to this path. Defaults to None.

        Returns:
            np.ndarray: A new image array with the prediction text overlaid.

        Raises:
            ValueError: If the input data is not a 3D image or if the
                predictions list is empty.
        """
        if input_data.ndim != 3:
            raise ValueError(f"Expected 3D image, got shape: {input_data.shape}")

        if not predictions:
            raise ValueError("Predictions list cannot be empty")

        vis_img = input_data.copy()
        vis_config = self.config.visualization
        font_scale = vis_config.font_scale
        thickness = vis_config.text_thickness if vis_config.text_thickness is not None else 1
        color = str_to_bgr(vis_config.text_color)
        top_k = self.config.params.get("top_k", 5)

        font = cv2.FONT_HERSHEY_SIMPLEX
        y_offset = 30

        for species, conf in predictions[:top_k]:
            text = f"{species}: {conf:.3f}"
            cv2.putText(
                vis_img,
                text,
                (10, y_offset),
                font,
                font_scale,
                color,
                thickness,
            )
            y_offset += 35

        if save_path:
            save_path = Path(save_path)
            save_path.parent.mkdir(parents=True, exist_ok=True)
            save_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
            cv2.imwrite(str(save_path), save_img)

        return vis_img

    def _evaluate_from_prediction(
        self,
        prediction: ClassificationPredictionType,
        ground_truth: ClassificationGroundTruthType,
    ) -> dict[str, float]:
        """Calculates core evaluation metrics for a single prediction.

        Args:
            prediction (ClassificationPredictionType): The model's prediction, which is
                a list of (species, confidence) tuples.
            ground_truth (ClassificationGroundTruthType): The true species name as a string.

        Returns:
            dict[str, float]: A dictionary of metrics including accuracy,
            confidence, top-1 correctness, and top-5 correctness.
        """
        if not prediction:
            return {
                "accuracy": 0.0,
                "confidence": 0.0,
                "top_1_correct": 0.0,
                "top_5_correct": 0.0,
            }

        pred_species = prediction[0][0]
        confidence = prediction[0][1]
        top_1_correct = float(pred_species == ground_truth)

        top_5_species = [p[0] for p in prediction[:5]]
        top_5_correct = float(ground_truth in top_5_species)

        return {
            "accuracy": top_1_correct,
            "confidence": confidence,
            "top_1_correct": top_1_correct,
            "top_5_correct": top_5_correct,
        }

    def _finalize_evaluation_report(
        self,
        aggregated_metrics: dict[str, float],
        predictions: list[ClassificationPredictionType],
        ground_truths: list[ClassificationGroundTruthType],
    ) -> dict[str, Any]:
        """Calculates and adds confusion matrix and ROC-AUC to the final report.

        Args:
            aggregated_metrics (dict[str, float]): A dictionary of metrics that have
                already been aggregated over the dataset.
            predictions (list[ClassificationPredictionType]): The list of all predictions.
            ground_truths (list[ClassificationGroundTruthType]): The list of all ground truths.

        Returns:
            dict[str, Any]: The updated report with the confusion matrix
            (as a list of lists) and the overall ROC-AUC score.
        """
        species_to_idx = {v: k for k, v in self.species_map.items()}
        class_labels = list(range(self.num_classes))
        y_true_indices, y_pred_indices, y_scores = [], [], []

        for gt, pred_list in zip(ground_truths, predictions):
            gt_str = gt
            if isinstance(gt_str, np.ndarray):
                if gt_str.shape == ():  # scalar array
                    gt_str = str(gt_str.item())
                else:
                    gt_str = str(gt_str.tolist())
            else:
                gt_str = str(gt_str)
            # Now safe to use as dict key
            if gt_str in species_to_idx and pred_list:
                true_idx = species_to_idx[gt_str]
                pred_str = pred_list[0][0]
                pred_idx = species_to_idx.get(pred_str, -1)

                y_true_indices.append(true_idx)
                y_pred_indices.append(pred_idx)

                prob_vector = [0.0] * self.num_classes
                for species, conf in pred_list:
                    class_idx = species_to_idx.get(species)
                    if class_idx is not None:
                        prob_vector[class_idx] = conf
                y_scores.append(prob_vector)

        if y_true_indices and y_pred_indices:
            valid_indices = [i for i, p_idx in enumerate(y_pred_indices) if p_idx != -1]
            if valid_indices:
                cm_y_true = [y_true_indices[i] for i in valid_indices]
                cm_y_pred = [y_pred_indices[i] for i in valid_indices]
                conf_matrix = confusion_matrix(
                    cm_y_true,
                    cm_y_pred,
                    labels=class_labels,
                )
                aggregated_metrics["confusion_matrix"] = conf_matrix.tolist()

        if y_scores and y_true_indices and len(np.unique(y_true_indices)) > 1:
            y_true_binarized = label_binarize(y_true_indices, classes=class_labels)
            try:
                roc_auc = roc_auc_score(
                    y_true_binarized,
                    np.array(y_scores),
                    multi_class="ovr",
                )
                aggregated_metrics["roc_auc"] = roc_auc
            except ValueError as e:
                self._logger.warning(f"Could not compute ROC AUC score: {e}")
                aggregated_metrics["roc_auc"] = 0.0

        return aggregated_metrics

    def _load_model(self) -> None:
        """Loads the pre-trained FastAI learner model from disk.

        This method uses the `set_posix_windows` context manager to ensure
        path compatibility across operating systems.

        Raises:
            RuntimeError: If the model file cannot be loaded, either because
                it is missing, corrupted, or dependencies are not met.
        """
        with set_posix_windows():
            try:
                self.learner = load_learner(self.model_path)
            except Exception as e:
                raise RuntimeError(
                    f"Failed to load existing model from {self.model_path}. "
                    f"Ensure the model file is valid and all dependencies are installed. Original error: {e}",
                ) from e

settings = settings instance-attribute

predictor_type = predictor_type instance-attribute

config: PredictorConfig property

Get the predictor configuration Pydantic model.

model_loaded: bool property

Check if the model is loaded.

model_path: Path property

Gets the path to the model weights file.

arch: str | None = self.config.model_arch instance-attribute

data_dir: Path = self.settings.dataset_dir instance-attribute

species_map: dict[int, str] = self.settings.species_config.species_map instance-attribute

num_classes: int = len(self.species_map) instance-attribute

__call__(input_data: np.ndarray, **kwargs: Any) -> Any

Convenience method that calls predict().

Source code in culicidaelab/core/base_predictor.py

def __call__(self, input_data: np.ndarray, **kwargs: Any) -> Any:
    """Convenience method that calls `predict()`."""
    if not self._model_loaded:
        self.load_model()
    return self.predict(input_data, **kwargs)

__enter__()

Context manager entry.

Source code in culicidaelab/core/base_predictor.py

def __enter__(self):
    """Context manager entry."""
    if not self._model_loaded:
        self.load_model()
    return self

__exit__(exc_type, exc_val, exc_tb)

Context manager exit.

Source code in culicidaelab/core/base_predictor.py

def __exit__(self, exc_type, exc_val, exc_tb):
    """Context manager exit."""
    pass

model_context()

A context manager for temporary model loading.

Ensures the model is loaded upon entering the context and unloaded upon exiting. This is useful for managing memory in pipelines.

Yields:

Name	Type	Description
BasePredictor		The predictor instance itself.

Example

with predictor.model_context(): ... predictions = predictor.predict(data)

Source code in culicidaelab/core/base_predictor.py

@contextmanager
def model_context(self):
    """A context manager for temporary model loading.

    Ensures the model is loaded upon entering the context and unloaded
    upon exiting. This is useful for managing memory in pipelines.

    Yields:
        BasePredictor: The predictor instance itself.

    Example:
        >>> with predictor.model_context():
        ...     predictions = predictor.predict(data)
    """
    was_loaded = self._model_loaded
    try:
        if not was_loaded:
            self.load_model()
        yield self
    finally:
        if not was_loaded and self._model_loaded:
            self.unload_model()

evaluate(ground_truth: GroundTruthType, prediction: PredictionType | None = None, input_data: np.ndarray | None = None, **predict_kwargs: Any) -> dict[str, float]

Evaluate a prediction against a ground truth.

Either prediction or input_data must be provided. If prediction is provided, it is used directly. If prediction is None, input_data is used to generate a new prediction.

Parameters:

Name	Type	Description	Default
ground_truth	GroundTruthType	The ground truth annotation.	required
prediction	PredictionType	A pre-computed prediction.	None
input_data	ndarray	Input data to generate a prediction from, if one isn't provided.	None
**predict_kwargs	Any	Additional arguments passed to the predict method.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing evaluation metrics for a
dict[str, float]	single item.

Raises:

Type	Description
ValueError	If neither prediction nor input_data is provided.

Source code in culicidaelab/core/base_predictor.py

def evaluate(
    self,
    ground_truth: GroundTruthType,
    prediction: PredictionType | None = None,
    input_data: np.ndarray | None = None,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate a prediction against a ground truth.

    Either `prediction` or `input_data` must be provided. If `prediction`
    is provided, it is used directly. If `prediction` is None, `input_data`
    is used to generate a new prediction.

    Args:
        ground_truth (GroundTruthType): The ground truth annotation.
        prediction (PredictionType, optional): A pre-computed prediction.
        input_data (np.ndarray, optional): Input data to generate a
            prediction from, if one isn't provided.
        **predict_kwargs (Any): Additional arguments passed to the `predict`
            method.

    Returns:
        dict[str, float]: Dictionary containing evaluation metrics for a
        single item.

    Raises:
        ValueError: If neither `prediction` nor `input_data` is provided.
    """
    if prediction is None:
        if input_data is not None:
            prediction = self.predict(input_data, **predict_kwargs)
        else:
            raise ValueError(
                "Either 'prediction' or 'input_data' must be provided.",
            )
    return self._evaluate_from_prediction(
        prediction=prediction,
        ground_truth=ground_truth,
    )

evaluate_batch(ground_truth_batch: list[GroundTruthType], predictions_batch: list[PredictionType] | None = None, input_data_batch: list[np.ndarray] | None = None, num_workers: int = 4, show_progress: bool = True, **predict_kwargs: Any) -> dict[str, float]

Evaluate on a batch of items using parallel processing.

Either predictions_batch or input_data_batch must be provided.

Parameters:

Name	Type	Description	Default
ground_truth_batch	list[GroundTruthType]	List of corresponding ground truth annotations.	required
predictions_batch	list[PredictionType]	A pre-computed list of predictions.	None
input_data_batch	list[ndarray]	List of input data to generate predictions from.	None
num_workers	int	Number of parallel workers for calculating metrics.	4
show_progress	bool	Whether to show a progress bar.	True
**predict_kwargs	Any	Additional arguments passed to predict_batch.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing aggregated evaluation metrics.

Raises:

Type	Description
ValueError	If the number of predictions does not match the number of ground truths.

Source code in culicidaelab/core/base_predictor.py

def evaluate_batch(
    self,
    ground_truth_batch: list[GroundTruthType],
    predictions_batch: list[PredictionType] | None = None,
    input_data_batch: list[np.ndarray] | None = None,
    num_workers: int = 4,
    show_progress: bool = True,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate on a batch of items using parallel processing.

    Either `predictions_batch` or `input_data_batch` must be provided.

    Args:
        ground_truth_batch (list[GroundTruthType]): List of corresponding
            ground truth annotations.
        predictions_batch (list[PredictionType], optional): A pre-computed
            list of predictions.
        input_data_batch (list[np.ndarray], optional): List of input data
            to generate predictions from.
        num_workers (int): Number of parallel workers for calculating metrics.
        show_progress (bool): Whether to show a progress bar.
        **predict_kwargs (Any): Additional arguments passed to `predict_batch`.

    Returns:
        dict[str, float]: Dictionary containing aggregated evaluation metrics.

    Raises:
        ValueError: If the number of predictions does not match the number
            of ground truths.
    """
    if predictions_batch is None:
        if input_data_batch is not None:
            predictions_batch = self.predict_batch(
                input_data_batch,
                show_progress=show_progress,
                **predict_kwargs,
            )
        else:
            raise ValueError(
                "Either 'predictions_batch' or 'input_data_batch' must be provided.",
            )

    if len(predictions_batch) != len(ground_truth_batch):
        raise ValueError(
            f"Number of predictions ({len(predictions_batch)}) must match "
            f"number of ground truths ({len(ground_truth_batch)}).",
        )

    per_item_metrics = self._calculate_metrics_parallel(
        predictions_batch,
        ground_truth_batch,
        num_workers,
        show_progress,
    )
    aggregated_metrics = self._aggregate_metrics(per_item_metrics)
    final_report = self._finalize_evaluation_report(
        aggregated_metrics,
        predictions_batch,
        ground_truth_batch,
    )
    return final_report

get_model_info() -> dict[str, Any]

Gets information about the loaded model.

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: A dictionary containing details about the model, such
dict[str, Any]	as architecture, path, etc.

Source code in culicidaelab/core/base_predictor.py

def get_model_info(self) -> dict[str, Any]:
    """Gets information about the loaded model.

    Returns:
        dict[str, Any]: A dictionary containing details about the model, such
        as architecture, path, etc.
    """
    return {
        "predictor_type": self.predictor_type,
        "model_path": str(self._model_path),
        "model_loaded": self._model_loaded,
        "config": self.config.model_dump(),
    }

load_model() -> None

Loads the model if it is not already loaded.

This is a convenience wrapper around _load_model that prevents reloading.

Raises:

Type	Description
RuntimeError	If model loading fails.

Source code in culicidaelab/core/base_predictor.py

def load_model(self) -> None:
    """Loads the model if it is not already loaded.

    This is a convenience wrapper around `_load_model` that prevents
    reloading.

    Raises:
        RuntimeError: If model loading fails.
    """
    if self._model_loaded:
        self._logger.info(f"Model for {self.predictor_type} already loaded")
        return

    try:
        self._logger.info(
            f"Loading model for {self.predictor_type} from {self._model_path}",
        )
        self._load_model()
        self._model_loaded = True
        self._logger.info(f"Successfully loaded model for {self.predictor_type}")
    except Exception as e:
        self._logger.error(f"Failed to load model for {self.predictor_type}: {e}")
        raise RuntimeError(
            f"Failed to load model for {self.predictor_type}: {e}",
        ) from e

unload_model() -> None

Unloads the model to free memory.

Source code in culicidaelab/core/base_predictor.py

def unload_model(self) -> None:
    """Unloads the model to free memory."""
    if self._model_loaded:
        self._model = None
        self._model_loaded = False
        self._logger.info(f"Unloaded model for {self.predictor_type}")

__init__(settings: Settings, load_model: bool = False) -> None

Initializes the MosquitoClassifier.

Source code in culicidaelab/predictors/classifier.py

def __init__(self, settings: Settings, load_model: bool = False) -> None:
    """Initializes the MosquitoClassifier."""
    provider_service = ProviderService(settings)
    weights_manager = ModelWeightsManager(
        settings=settings,
        provider_service=provider_service,
    )
    super().__init__(
        settings=settings,
        predictor_type="classifier",
        weights_manager=weights_manager,
        load_model=load_model,
    )

    self.arch: str | None = self.config.model_arch
    self.data_dir: Path = self.settings.dataset_dir
    self.species_map: dict[int, str] = self.settings.species_config.species_map
    self.num_classes: int = len(self.species_map)

get_class_index(species_name: str) -> int | None

Retrieves the class index for a given species name.

Parameters:

Name	Type	Description	Default
species_name	str	The name of the species.	required

Returns:

Type	Description
int | None	int | None: The corresponding class index if found, otherwise None.

Source code in culicidaelab/predictors/classifier.py

def get_class_index(self, species_name: str) -> int | None:
    """Retrieves the class index for a given species name.

    Args:
        species_name (str): The name of the species.

    Returns:
        int | None: The corresponding class index if found, otherwise None.
    """
    return self.settings.species_config.get_index_by_species(species_name)

get_species_names() -> list[str]

Gets a sorted list of all species names known to the classifier.

The list is ordered by the class index.

Returns:

Type	Description
list[str]	list[str]: A list of species names.

Source code in culicidaelab/predictors/classifier.py

def get_species_names(self) -> list[str]:
    """Gets a sorted list of all species names known to the classifier.

    The list is ordered by the class index.

    Returns:
        list[str]: A list of species names.
    """
    return [self.species_map[i] for i in sorted(self.species_map.keys())]

predict(input_data: np.ndarray, **kwargs: Any) -> ClassificationPredictionType

Classifies the mosquito species in a single image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	An input image as a NumPy array with a shape of (H, W, 3) in RGB format. Values can be uint8 [0, 255] or float [0, 1].	required
**kwargs	Any	Additional arguments (not used).	{}

Returns:

Name	Type	Description
ClassificationPredictionType	ClassificationPredictionType	A list of (species_name, confidence)
	ClassificationPredictionType	tuples, sorted in descending order of confidence.

Raises:

Type	Description
RuntimeError	If the model has not been loaded.
ValueError	If the input data has an invalid shape, data type, or is not a valid image.

Source code in culicidaelab/predictors/classifier.py

def predict(
    self,
    input_data: np.ndarray,
    **kwargs: Any,
) -> ClassificationPredictionType:
    """Classifies the mosquito species in a single image.

    Args:
        input_data (np.ndarray): An input image as a NumPy array with a
            shape of (H, W, 3) in RGB format. Values can be uint8
            [0, 255] or float [0, 1].
        **kwargs (Any): Additional arguments (not used).

    Returns:
        ClassificationPredictionType: A list of (species_name, confidence)
        tuples, sorted in descending order of confidence.

    Raises:
        RuntimeError: If the model has not been loaded.
        ValueError: If the input data has an invalid shape, data type,
            or is not a valid image.
    """
    if not self.model_loaded:
        raise RuntimeError(
            "Model is not loaded. Call load_model() or use a context manager.",
        )

    if not isinstance(input_data, np.ndarray):
        input_data = np.array(input_data)

    if input_data.ndim != 3 or input_data.shape[2] != 3:
        raise ValueError(f"Expected 3D RGB image, got shape: {input_data.shape}")

    if input_data.dtype == np.uint8:
        image = Image.fromarray(input_data)
    elif input_data.dtype in [np.float32, np.float64]:
        image = Image.fromarray((input_data * 255).astype(np.uint8))
    else:
        raise ValueError(f"Unsupported dtype: {input_data.dtype}")

    with set_posix_windows():
        _, _, probabilities = self.learner.predict(image)

    species_probs = []
    for idx, prob in enumerate(probabilities):
        species_name = self.species_map.get(idx, f"unknown_{idx}")
        species_probs.append((species_name, float(prob)))

    species_probs.sort(key=lambda x: x[1], reverse=True)
    return species_probs

predict_batch(input_data_batch: list[Any], show_progress: bool = False, **kwargs: Any) -> list[ClassificationPredictionType]

Classifies mosquito species in a batch of images.

Note: This method currently iterates and calls predict for each image. True batch processing is not yet implemented.

Parameters:

Name	Type	Description	Default
input_data_batch	list[Any]	A list of input images, where each image is a NumPy array.	required
show_progress	bool	If True, a progress bar is displayed. Defaults to False.	False
**kwargs	Any	Additional arguments passed to predict.	{}

Returns:

Type	Description
list[ClassificationPredictionType]	list[ClassificationPredictionType]: A list of prediction results,
list[ClassificationPredictionType]	where each result corresponds to an input image.

Source code in culicidaelab/predictors/classifier.py

def predict_batch(
    self,
    input_data_batch: list[Any],
    show_progress: bool = False,
    **kwargs: Any,
) -> list[ClassificationPredictionType]:
    """Classifies mosquito species in a batch of images.

    Note: This method currently iterates and calls `predict` for each image.
    True batch processing is not yet implemented.

    Args:
        input_data_batch (list[Any]): A list of input images, where each
            image is a NumPy array.
        show_progress (bool, optional): If True, a progress bar is displayed.
            Defaults to False.
        **kwargs (Any): Additional arguments passed to `predict`.

    Returns:
        list[ClassificationPredictionType]: A list of prediction results,
        where each result corresponds to an input image.
    """
    results = []
    for img in input_data_batch:
        results.append(self.predict(img, **kwargs))
    return results

visualize(input_data: np.ndarray, predictions: ClassificationPredictionType, save_path: str | Path | None = None) -> np.ndarray

Overlays classification results on an image.

This method draws the top-k predictions and their confidence scores onto the input image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	The original image (H, W, 3) as a NumPy array.	required
predictions	ClassificationPredictionType	The prediction output from the predict method.	required
save_path	str | Path | None	If provided, the visualized image will be saved to this path. Defaults to None.	None

Returns:

Type	Description
ndarray	np.ndarray: A new image array with the prediction text overlaid.

Raises:

Type	Description
ValueError	If the input data is not a 3D image or if the predictions list is empty.

Source code in culicidaelab/predictors/classifier.py

def visualize(
    self,
    input_data: np.ndarray,
    predictions: ClassificationPredictionType,
    save_path: str | Path | None = None,
) -> np.ndarray:
    """Overlays classification results on an image.

    This method draws the top-k predictions and their confidence scores
    onto the input image.

    Args:
        input_data (np.ndarray): The original image (H, W, 3) as a NumPy array.
        predictions (ClassificationPredictionType): The prediction output from
            the `predict` method.
        save_path (str | Path | None, optional): If provided, the visualized
            image will be saved to this path. Defaults to None.

    Returns:
        np.ndarray: A new image array with the prediction text overlaid.

    Raises:
        ValueError: If the input data is not a 3D image or if the
            predictions list is empty.
    """
    if input_data.ndim != 3:
        raise ValueError(f"Expected 3D image, got shape: {input_data.shape}")

    if not predictions:
        raise ValueError("Predictions list cannot be empty")

    vis_img = input_data.copy()
    vis_config = self.config.visualization
    font_scale = vis_config.font_scale
    thickness = vis_config.text_thickness if vis_config.text_thickness is not None else 1
    color = str_to_bgr(vis_config.text_color)
    top_k = self.config.params.get("top_k", 5)

    font = cv2.FONT_HERSHEY_SIMPLEX
    y_offset = 30

    for species, conf in predictions[:top_k]:
        text = f"{species}: {conf:.3f}"
        cv2.putText(
            vis_img,
            text,
            (10, y_offset),
            font,
            font_scale,
            color,
            thickness,
        )
        y_offset += 35

    if save_path:
        save_path = Path(save_path)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        save_img = cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(str(save_path), save_img)

    return vis_img

MosquitoDetector

Detects mosquitos in images using a YOLO model.

This class loads a YOLO model and provides methods for predicting bounding boxes on single or batches of images, visualizing results, and evaluating detection performance against ground truth data.

Parameters:

Name	Type	Description	Default
settings	Settings	The main settings object for the library.	required
load_model	bool	If True, the model is loaded upon initialization. Defaults to False.	False

Attributes:

Name	Type	Description
confidence_threshold	float	The minimum confidence score for a detection to be considered valid.
iou_threshold	float	The IoU threshold for non-maximum suppression.
max_detections	int	The maximum number of detections to return per image.

Source code in culicidaelab/predictors/detector.py

class MosquitoDetector(BasePredictor[DetectionPredictionType, DetectionGroundTruthType]):
    """Detects mosquitos in images using a YOLO model.

    This class loads a YOLO model and provides methods for predicting bounding
    boxes on single or batches of images, visualizing results, and evaluating
    detection performance against ground truth data.

    Args:
        settings (Settings): The main settings object for the library.
        load_model (bool, optional): If True, the model is loaded upon
            initialization. Defaults to False.

    Attributes:
        confidence_threshold (float): The minimum confidence score for a
            detection to be considered valid.
        iou_threshold (float): The IoU threshold for non-maximum suppression.
        max_detections (int): The maximum number of detections to return per image.
    """

    def __init__(self, settings: Settings, load_model: bool = False) -> None:
        """Initializes the MosquitoDetector."""
        provider_service = ProviderService(settings)
        weights_manager = ModelWeightsManager(
            settings=settings,
            provider_service=provider_service,
        )
        super().__init__(
            settings=settings,
            predictor_type="detector",
            weights_manager=weights_manager,
            load_model=load_model,
        )
        self.confidence_threshold: float = self.config.confidence or 0.5
        self.iou_threshold: float = self.config.params.get("iou_threshold", 0.45)
        self.max_detections: int = self.config.params.get("max_detections", 300)

    def predict(self, input_data: np.ndarray, **kwargs: Any) -> DetectionPredictionType:
        """Detects mosquitos in a single image.

        Args:
            input_data (np.ndarray): The input image as a NumPy array.
            **kwargs (Any): Optional keyword arguments, including:
                confidence_threshold (float): Override the default confidence
                    threshold for this prediction.

        Returns:
            DetectionPredictionType: A list of detection tuples. Each tuple is
            (center_x, center_y, width, height, confidence). Returns an empty
            list if no mosquitos are found.

        Raises:
            RuntimeError: If the model fails to load or if prediction fails.
        """
        if not self.model_loaded or self._model is None:
            self.load_model()
            if self._model is None:
                raise RuntimeError("Failed to load model")

        confidence_threshold = kwargs.get("confidence_threshold", self.confidence_threshold)

        try:
            results = self._model(
                source=input_data,
                conf=confidence_threshold,
                iou=self.iou_threshold,
                max_det=self.max_detections,
                verbose=False,
            )
        except Exception as e:
            logger.error(f"Prediction failed: {e}", exc_info=True)
            raise RuntimeError(f"Prediction failed: {e}") from e

        detections: DetectionPredictionType = []
        if results:
            boxes = results[0].boxes
            for box in boxes:
                xyxy_tensor = box.xyxy[0]
                x1, y1, x2, y2 = xyxy_tensor.cpu().numpy()
                conf = float(box.conf[0])
                w, h = x2 - x1, y2 - y1
                center_x, center_y = x1 + w / 2, y1 + h / 2
                detections.append((center_x, center_y, w, h, conf))
        return detections

    def predict_batch(
        self,
        input_data_batch: list[np.ndarray],
        show_progress: bool = True,
        **kwargs: Any,
    ) -> list[DetectionPredictionType]:
        """Detects mosquitos in a batch of images using YOLO's native batching.

        Args:
            input_data_batch (list[np.ndarray]): A list of input images.
            show_progress (bool, optional): If True, a progress bar is shown.
                Defaults to True.
            **kwargs (Any): Additional arguments (not used).

        Returns:
            list[DetectionPredictionType]: A list where each item is the list
            of detections for the corresponding image in the input batch.

        Raises:
            RuntimeError: If the model is not loaded.
        """
        if not self.model_loaded or self._model is None:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        # The tqdm iterator is just for visual feedback; actual prediction happens in one go.
        iterator = tqdm(range(len(input_data_batch)), desc="Predicting detection batch", disable=not show_progress)

        yolo_results = self._model(
            source=input_data_batch,
            conf=self.confidence_threshold,
            iou=self.iou_threshold,
            max_det=self.max_detections,
            stream=False,
            verbose=False,
        )

        all_predictions: list[DetectionPredictionType] = []
        for r in yolo_results:
            detections = []
            for box in r.boxes:
                xyxy_tensor = box.xyxy[0]
                x1, y1, x2, y2 = xyxy_tensor.cpu().numpy()
                w, h = x2 - x1, y2 - y1
                center_x, center_y = x1 + w / 2, y1 + h / 2
                conf = float(box.conf[0])
                detections.append((center_x, center_y, w, h, conf))
            all_predictions.append(detections)
            iterator.update(1)
        iterator.close()

        return all_predictions

    def visualize(
        self,
        input_data: np.ndarray,
        predictions: DetectionPredictionType,
        save_path: str | Path | None = None,
    ) -> np.ndarray:
        """Draws predicted bounding boxes on an image.

        Args:
            input_data (np.ndarray): The original image.
            predictions (DetectionPredictionType): The list of detections from `predict`.
            save_path (str | Path | None, optional): If provided, the output
                image is saved to this path. Defaults to None.

        Returns:
            np.ndarray: A new image array with bounding boxes and confidence
            scores drawn on it.
        """
        vis_img = input_data.copy()
        vis_config = self.config.visualization
        box_color = str_to_bgr(vis_config.box_color)
        text_color = str_to_bgr(vis_config.text_color)
        font_scale = vis_config.font_scale
        thickness = vis_config.box_thickness

        for x, y, w, h, conf in predictions:
            x1, y1 = int(x - w / 2), int(y - h / 2)
            x2, y2 = int(x + w / 2), int(y + h / 2)
            cv2.rectangle(vis_img, (x1, y1), (x2, y2), box_color, thickness)
            text = f"{conf:.2f}"
            cv2.putText(vis_img, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color, thickness)

        if save_path:
            cv2.imwrite(str(save_path), cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR))

        return vis_img

    def _calculate_iou(self, box1_xywh: tuple, box2_xywh: tuple) -> float:
        """Calculates Intersection over Union (IoU) for two boxes.

        Args:
            box1_xywh (tuple): The first box in (cx, cy, w, h) format.
            box2_xywh (tuple): The second box in (cx, cy, w, h) format.

        Returns:
            float: The IoU score between 0.0 and 1.0.
        """
        b1_x1, b1_y1 = box1_xywh[0] - box1_xywh[2] / 2, box1_xywh[1] - box1_xywh[3] / 2
        b1_x2, b1_y2 = box1_xywh[0] + box1_xywh[2] / 2, box1_xywh[1] + box1_xywh[3] / 2
        b2_x1, b2_y1 = box2_xywh[0] - box2_xywh[2] / 2, box2_xywh[1] - box2_xywh[3] / 2
        b2_x2, b2_y2 = box2_xywh[0] + box2_xywh[2] / 2, box2_xywh[1] + box2_xywh[3] / 2

        inter_x1, inter_y1 = max(b1_x1, b2_x1), max(b1_y1, b2_y1)
        inter_x2, inter_y2 = min(b1_x2, b2_x2), min(b1_y2, b2_y2)
        intersection = max(0, inter_x2 - inter_x1) * max(0, inter_y2 - inter_y1)

        area1, area2 = box1_xywh[2] * box1_xywh[3], box2_xywh[2] * box2_xywh[3]
        union = area1 + area2 - intersection
        return float(intersection / union) if union > 0 else 0.0

    def _evaluate_from_prediction(
        self,
        prediction: DetectionPredictionType,
        ground_truth: DetectionGroundTruthType,
    ) -> dict[str, float]:
        """Calculates detection metrics for a single image's predictions.

        This computes precision, recall, F1-score, Average Precision (AP),
        and mean IoU for a set of predicted boxes against ground truth boxes.

        Args:
            prediction (DetectionPredictionType): A list of predicted boxes with
                confidence scores: `[(x, y, w, h, conf), ...]`.
            ground_truth (DetectionGroundTruthType): A list of ground truth
                boxes: `[(x, y, w, h), ...]`.

        Returns:
            dict[str, float]: A dictionary containing the calculated metrics.
        """
        if not ground_truth and not prediction:
            return {"precision": 1.0, "recall": 1.0, "f1": 1.0, "ap": 1.0, "mean_iou": 0.0}
        if not ground_truth:  # False positives exist
            return {"precision": 0.0, "recall": 0.0, "f1": 0.0, "ap": 0.0, "mean_iou": 0.0}
        if not prediction:  # False negatives exist
            return {"precision": 0.0, "recall": 0.0, "f1": 0.0, "ap": 0.0, "mean_iou": 0.0}

        predictions_sorted = sorted(prediction, key=lambda x: x[4], reverse=True)
        tp = np.zeros(len(predictions_sorted))
        fp = np.zeros(len(predictions_sorted))
        gt_matched = [False] * len(ground_truth)
        all_ious_for_mean = []
        iou_threshold = self.iou_threshold

        for i, pred_box_with_conf in enumerate(predictions_sorted):
            pred_box = pred_box_with_conf[:4]
            best_iou, best_gt_idx = 0.0, -1

            for j, gt_box in enumerate(ground_truth):
                if not gt_matched[j]:
                    iou = self._calculate_iou(pred_box, gt_box)
                    if iou > best_iou:
                        best_iou = iou
                        best_gt_idx = j

            if best_gt_idx != -1:
                all_ious_for_mean.append(best_iou)

            if best_iou >= iou_threshold:
                if not gt_matched[best_gt_idx]:
                    tp[i] = 1
                    gt_matched[best_gt_idx] = True
                else:  # Matched a GT box that was already matched
                    fp[i] = 1
            else:
                fp[i] = 1

        mean_iou_val = float(np.mean(all_ious_for_mean)) if all_ious_for_mean else 0.0
        fp_cumsum, tp_cumsum = np.cumsum(fp), np.cumsum(tp)
        recall_curve = tp_cumsum / len(ground_truth)
        precision_curve = tp_cumsum / (tp_cumsum + fp_cumsum + 1e-9)

        ap = 0.0
        for t in np.linspace(0, 1, 11):  # 11-point interpolation
            precisions_at_recall_t = precision_curve[recall_curve >= t]
            ap += np.max(precisions_at_recall_t) if len(precisions_at_recall_t) > 0 else 0.0
        ap /= 11.0

        final_precision = precision_curve[-1] if len(precision_curve) > 0 else 0.0
        final_recall = recall_curve[-1] if len(recall_curve) > 0 else 0.0
        f1 = (
            2 * (final_precision * final_recall) / (final_precision + final_recall + 1e-9)
            if (final_precision + final_recall) > 0
            else 0.0
        )

        return {
            "precision": float(final_precision),
            "recall": float(final_recall),
            "f1": float(f1),
            "ap": float(ap),
            "mean_iou": mean_iou_val,
        }

    def _load_model(self) -> None:
        """Loads the YOLO object detection model from the specified path.

        Raises:
            RuntimeError: If the model cannot be loaded from the path
                specified in the configuration.
        """
        try:
            logger.info(f"Loading YOLO model from: {self.model_path}")
            self._model = YOLO(str(self.model_path), task="detect")

            if self._model and hasattr(self.config, "device") and self.config.device:
                device = str(self.config.device)
                logger.info(f"Moving model to device: {device}")
                self._model.to(device)

            logger.info("YOLO model loaded successfully.")
        except Exception as e:
            logger.error(f"Failed to load YOLO model: {e}", exc_info=True)
            self._model = None
            raise RuntimeError(f"Could not load YOLO model from {self.model_path}.") from e

settings = settings instance-attribute

predictor_type = predictor_type instance-attribute

config: PredictorConfig property

Get the predictor configuration Pydantic model.

model_loaded: bool property

Check if the model is loaded.

model_path: Path property

Gets the path to the model weights file.

confidence_threshold: float = self.config.confidence or 0.5 instance-attribute

iou_threshold: float = self.config.params.get('iou_threshold', 0.45) instance-attribute

max_detections: int = self.config.params.get('max_detections', 300) instance-attribute

__call__(input_data: np.ndarray, **kwargs: Any) -> Any

Convenience method that calls predict().

Source code in culicidaelab/core/base_predictor.py

def __call__(self, input_data: np.ndarray, **kwargs: Any) -> Any:
    """Convenience method that calls `predict()`."""
    if not self._model_loaded:
        self.load_model()
    return self.predict(input_data, **kwargs)

__enter__()

Context manager entry.

Source code in culicidaelab/core/base_predictor.py

def __enter__(self):
    """Context manager entry."""
    if not self._model_loaded:
        self.load_model()
    return self

__exit__(exc_type, exc_val, exc_tb)

Context manager exit.

Source code in culicidaelab/core/base_predictor.py

def __exit__(self, exc_type, exc_val, exc_tb):
    """Context manager exit."""
    pass

model_context()

A context manager for temporary model loading.

Ensures the model is loaded upon entering the context and unloaded upon exiting. This is useful for managing memory in pipelines.

Yields:

Name	Type	Description
BasePredictor		The predictor instance itself.

Example

with predictor.model_context(): ... predictions = predictor.predict(data)

Source code in culicidaelab/core/base_predictor.py

@contextmanager
def model_context(self):
    """A context manager for temporary model loading.

    Ensures the model is loaded upon entering the context and unloaded
    upon exiting. This is useful for managing memory in pipelines.

    Yields:
        BasePredictor: The predictor instance itself.

    Example:
        >>> with predictor.model_context():
        ...     predictions = predictor.predict(data)
    """
    was_loaded = self._model_loaded
    try:
        if not was_loaded:
            self.load_model()
        yield self
    finally:
        if not was_loaded and self._model_loaded:
            self.unload_model()

evaluate(ground_truth: GroundTruthType, prediction: PredictionType | None = None, input_data: np.ndarray | None = None, **predict_kwargs: Any) -> dict[str, float]

Evaluate a prediction against a ground truth.

Either prediction or input_data must be provided. If prediction is provided, it is used directly. If prediction is None, input_data is used to generate a new prediction.

Parameters:

Name	Type	Description	Default
ground_truth	GroundTruthType	The ground truth annotation.	required
prediction	PredictionType	A pre-computed prediction.	None
input_data	ndarray	Input data to generate a prediction from, if one isn't provided.	None
**predict_kwargs	Any	Additional arguments passed to the predict method.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing evaluation metrics for a
dict[str, float]	single item.

Raises:

Type	Description
ValueError	If neither prediction nor input_data is provided.

Source code in culicidaelab/core/base_predictor.py

def evaluate(
    self,
    ground_truth: GroundTruthType,
    prediction: PredictionType | None = None,
    input_data: np.ndarray | None = None,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate a prediction against a ground truth.

    Either `prediction` or `input_data` must be provided. If `prediction`
    is provided, it is used directly. If `prediction` is None, `input_data`
    is used to generate a new prediction.

    Args:
        ground_truth (GroundTruthType): The ground truth annotation.
        prediction (PredictionType, optional): A pre-computed prediction.
        input_data (np.ndarray, optional): Input data to generate a
            prediction from, if one isn't provided.
        **predict_kwargs (Any): Additional arguments passed to the `predict`
            method.

    Returns:
        dict[str, float]: Dictionary containing evaluation metrics for a
        single item.

    Raises:
        ValueError: If neither `prediction` nor `input_data` is provided.
    """
    if prediction is None:
        if input_data is not None:
            prediction = self.predict(input_data, **predict_kwargs)
        else:
            raise ValueError(
                "Either 'prediction' or 'input_data' must be provided.",
            )
    return self._evaluate_from_prediction(
        prediction=prediction,
        ground_truth=ground_truth,
    )

evaluate_batch(ground_truth_batch: list[GroundTruthType], predictions_batch: list[PredictionType] | None = None, input_data_batch: list[np.ndarray] | None = None, num_workers: int = 4, show_progress: bool = True, **predict_kwargs: Any) -> dict[str, float]

Evaluate on a batch of items using parallel processing.

Either predictions_batch or input_data_batch must be provided.

Parameters:

Name	Type	Description	Default
ground_truth_batch	list[GroundTruthType]	List of corresponding ground truth annotations.	required
predictions_batch	list[PredictionType]	A pre-computed list of predictions.	None
input_data_batch	list[ndarray]	List of input data to generate predictions from.	None
num_workers	int	Number of parallel workers for calculating metrics.	4
show_progress	bool	Whether to show a progress bar.	True
**predict_kwargs	Any	Additional arguments passed to predict_batch.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing aggregated evaluation metrics.

Raises:

Type	Description
ValueError	If the number of predictions does not match the number of ground truths.

Source code in culicidaelab/core/base_predictor.py

def evaluate_batch(
    self,
    ground_truth_batch: list[GroundTruthType],
    predictions_batch: list[PredictionType] | None = None,
    input_data_batch: list[np.ndarray] | None = None,
    num_workers: int = 4,
    show_progress: bool = True,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate on a batch of items using parallel processing.

    Either `predictions_batch` or `input_data_batch` must be provided.

    Args:
        ground_truth_batch (list[GroundTruthType]): List of corresponding
            ground truth annotations.
        predictions_batch (list[PredictionType], optional): A pre-computed
            list of predictions.
        input_data_batch (list[np.ndarray], optional): List of input data
            to generate predictions from.
        num_workers (int): Number of parallel workers for calculating metrics.
        show_progress (bool): Whether to show a progress bar.
        **predict_kwargs (Any): Additional arguments passed to `predict_batch`.

    Returns:
        dict[str, float]: Dictionary containing aggregated evaluation metrics.

    Raises:
        ValueError: If the number of predictions does not match the number
            of ground truths.
    """
    if predictions_batch is None:
        if input_data_batch is not None:
            predictions_batch = self.predict_batch(
                input_data_batch,
                show_progress=show_progress,
                **predict_kwargs,
            )
        else:
            raise ValueError(
                "Either 'predictions_batch' or 'input_data_batch' must be provided.",
            )

    if len(predictions_batch) != len(ground_truth_batch):
        raise ValueError(
            f"Number of predictions ({len(predictions_batch)}) must match "
            f"number of ground truths ({len(ground_truth_batch)}).",
        )

    per_item_metrics = self._calculate_metrics_parallel(
        predictions_batch,
        ground_truth_batch,
        num_workers,
        show_progress,
    )
    aggregated_metrics = self._aggregate_metrics(per_item_metrics)
    final_report = self._finalize_evaluation_report(
        aggregated_metrics,
        predictions_batch,
        ground_truth_batch,
    )
    return final_report

get_model_info() -> dict[str, Any]

Gets information about the loaded model.

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: A dictionary containing details about the model, such
dict[str, Any]	as architecture, path, etc.

Source code in culicidaelab/core/base_predictor.py

def get_model_info(self) -> dict[str, Any]:
    """Gets information about the loaded model.

    Returns:
        dict[str, Any]: A dictionary containing details about the model, such
        as architecture, path, etc.
    """
    return {
        "predictor_type": self.predictor_type,
        "model_path": str(self._model_path),
        "model_loaded": self._model_loaded,
        "config": self.config.model_dump(),
    }

load_model() -> None

Loads the model if it is not already loaded.

This is a convenience wrapper around _load_model that prevents reloading.

Raises:

Type	Description
RuntimeError	If model loading fails.

Source code in culicidaelab/core/base_predictor.py

def load_model(self) -> None:
    """Loads the model if it is not already loaded.

    This is a convenience wrapper around `_load_model` that prevents
    reloading.

    Raises:
        RuntimeError: If model loading fails.
    """
    if self._model_loaded:
        self._logger.info(f"Model for {self.predictor_type} already loaded")
        return

    try:
        self._logger.info(
            f"Loading model for {self.predictor_type} from {self._model_path}",
        )
        self._load_model()
        self._model_loaded = True
        self._logger.info(f"Successfully loaded model for {self.predictor_type}")
    except Exception as e:
        self._logger.error(f"Failed to load model for {self.predictor_type}: {e}")
        raise RuntimeError(
            f"Failed to load model for {self.predictor_type}: {e}",
        ) from e

unload_model() -> None

Unloads the model to free memory.

Source code in culicidaelab/core/base_predictor.py

def unload_model(self) -> None:
    """Unloads the model to free memory."""
    if self._model_loaded:
        self._model = None
        self._model_loaded = False
        self._logger.info(f"Unloaded model for {self.predictor_type}")

__init__(settings: Settings, load_model: bool = False) -> None

Initializes the MosquitoDetector.

Source code in culicidaelab/predictors/detector.py

def __init__(self, settings: Settings, load_model: bool = False) -> None:
    """Initializes the MosquitoDetector."""
    provider_service = ProviderService(settings)
    weights_manager = ModelWeightsManager(
        settings=settings,
        provider_service=provider_service,
    )
    super().__init__(
        settings=settings,
        predictor_type="detector",
        weights_manager=weights_manager,
        load_model=load_model,
    )
    self.confidence_threshold: float = self.config.confidence or 0.5
    self.iou_threshold: float = self.config.params.get("iou_threshold", 0.45)
    self.max_detections: int = self.config.params.get("max_detections", 300)

predict(input_data: np.ndarray, **kwargs: Any) -> DetectionPredictionType

Detects mosquitos in a single image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	The input image as a NumPy array.	required
**kwargs	Any	Optional keyword arguments, including: confidence_threshold (float): Override the default confidence threshold for this prediction.	{}

Returns:

Name	Type	Description
DetectionPredictionType	DetectionPredictionType	A list of detection tuples. Each tuple is
	DetectionPredictionType	(center_x, center_y, width, height, confidence). Returns an empty
	DetectionPredictionType	list if no mosquitos are found.

Raises:

Type	Description
RuntimeError	If the model fails to load or if prediction fails.

Source code in culicidaelab/predictors/detector.py

def predict(self, input_data: np.ndarray, **kwargs: Any) -> DetectionPredictionType:
    """Detects mosquitos in a single image.

    Args:
        input_data (np.ndarray): The input image as a NumPy array.
        **kwargs (Any): Optional keyword arguments, including:
            confidence_threshold (float): Override the default confidence
                threshold for this prediction.

    Returns:
        DetectionPredictionType: A list of detection tuples. Each tuple is
        (center_x, center_y, width, height, confidence). Returns an empty
        list if no mosquitos are found.

    Raises:
        RuntimeError: If the model fails to load or if prediction fails.
    """
    if not self.model_loaded or self._model is None:
        self.load_model()
        if self._model is None:
            raise RuntimeError("Failed to load model")

    confidence_threshold = kwargs.get("confidence_threshold", self.confidence_threshold)

    try:
        results = self._model(
            source=input_data,
            conf=confidence_threshold,
            iou=self.iou_threshold,
            max_det=self.max_detections,
            verbose=False,
        )
    except Exception as e:
        logger.error(f"Prediction failed: {e}", exc_info=True)
        raise RuntimeError(f"Prediction failed: {e}") from e

    detections: DetectionPredictionType = []
    if results:
        boxes = results[0].boxes
        for box in boxes:
            xyxy_tensor = box.xyxy[0]
            x1, y1, x2, y2 = xyxy_tensor.cpu().numpy()
            conf = float(box.conf[0])
            w, h = x2 - x1, y2 - y1
            center_x, center_y = x1 + w / 2, y1 + h / 2
            detections.append((center_x, center_y, w, h, conf))
    return detections

predict_batch(input_data_batch: list[np.ndarray], show_progress: bool = True, **kwargs: Any) -> list[DetectionPredictionType]

Detects mosquitos in a batch of images using YOLO's native batching.

Parameters:

Name	Type	Description	Default
input_data_batch	list[ndarray]	A list of input images.	required
show_progress	bool	If True, a progress bar is shown. Defaults to True.	True
**kwargs	Any	Additional arguments (not used).	{}

Returns:

Type	Description
list[DetectionPredictionType]	list[DetectionPredictionType]: A list where each item is the list
list[DetectionPredictionType]	of detections for the corresponding image in the input batch.

Raises:

Type	Description
RuntimeError	If the model is not loaded.

Source code in culicidaelab/predictors/detector.py

def predict_batch(
    self,
    input_data_batch: list[np.ndarray],
    show_progress: bool = True,
    **kwargs: Any,
) -> list[DetectionPredictionType]:
    """Detects mosquitos in a batch of images using YOLO's native batching.

    Args:
        input_data_batch (list[np.ndarray]): A list of input images.
        show_progress (bool, optional): If True, a progress bar is shown.
            Defaults to True.
        **kwargs (Any): Additional arguments (not used).

    Returns:
        list[DetectionPredictionType]: A list where each item is the list
        of detections for the corresponding image in the input batch.

    Raises:
        RuntimeError: If the model is not loaded.
    """
    if not self.model_loaded or self._model is None:
        raise RuntimeError("Model not loaded. Call load_model() first.")

    # The tqdm iterator is just for visual feedback; actual prediction happens in one go.
    iterator = tqdm(range(len(input_data_batch)), desc="Predicting detection batch", disable=not show_progress)

    yolo_results = self._model(
        source=input_data_batch,
        conf=self.confidence_threshold,
        iou=self.iou_threshold,
        max_det=self.max_detections,
        stream=False,
        verbose=False,
    )

    all_predictions: list[DetectionPredictionType] = []
    for r in yolo_results:
        detections = []
        for box in r.boxes:
            xyxy_tensor = box.xyxy[0]
            x1, y1, x2, y2 = xyxy_tensor.cpu().numpy()
            w, h = x2 - x1, y2 - y1
            center_x, center_y = x1 + w / 2, y1 + h / 2
            conf = float(box.conf[0])
            detections.append((center_x, center_y, w, h, conf))
        all_predictions.append(detections)
        iterator.update(1)
    iterator.close()

    return all_predictions

visualize(input_data: np.ndarray, predictions: DetectionPredictionType, save_path: str | Path | None = None) -> np.ndarray

Draws predicted bounding boxes on an image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	The original image.	required
predictions	DetectionPredictionType	The list of detections from predict.	required
save_path	str | Path | None	If provided, the output image is saved to this path. Defaults to None.	None

Returns:

Type	Description
ndarray	np.ndarray: A new image array with bounding boxes and confidence
ndarray	scores drawn on it.

Source code in culicidaelab/predictors/detector.py

def visualize(
    self,
    input_data: np.ndarray,
    predictions: DetectionPredictionType,
    save_path: str | Path | None = None,
) -> np.ndarray:
    """Draws predicted bounding boxes on an image.

    Args:
        input_data (np.ndarray): The original image.
        predictions (DetectionPredictionType): The list of detections from `predict`.
        save_path (str | Path | None, optional): If provided, the output
            image is saved to this path. Defaults to None.

    Returns:
        np.ndarray: A new image array with bounding boxes and confidence
        scores drawn on it.
    """
    vis_img = input_data.copy()
    vis_config = self.config.visualization
    box_color = str_to_bgr(vis_config.box_color)
    text_color = str_to_bgr(vis_config.text_color)
    font_scale = vis_config.font_scale
    thickness = vis_config.box_thickness

    for x, y, w, h, conf in predictions:
        x1, y1 = int(x - w / 2), int(y - h / 2)
        x2, y2 = int(x + w / 2), int(y + h / 2)
        cv2.rectangle(vis_img, (x1, y1), (x2, y2), box_color, thickness)
        text = f"{conf:.2f}"
        cv2.putText(vis_img, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color, thickness)

    if save_path:
        cv2.imwrite(str(save_path), cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR))

    return vis_img

MosquitoSegmenter

Segments mosquitos in images using the SAM2 model.

This class provides methods to load a SAM2 model, generate segmentation masks for entire images or specific regions defined by bounding boxes, and visualize the resulting masks.

Parameters:

Name	Type	Description	Default
settings	Settings	The main settings object for the library.	required
load_model	bool	If True, the model is loaded upon initialization. Defaults to False.	False

Source code in culicidaelab/predictors/segmenter.py

class MosquitoSegmenter(BasePredictor[SegmentationPredictionType, SegmentationGroundTruthType]):
    """Segments mosquitos in images using the SAM2 model.

    This class provides methods to load a SAM2 model, generate segmentation
    masks for entire images or specific regions defined by bounding boxes,
    and visualize the resulting masks.

    Args:
        settings (Settings): The main settings object for the library.
        load_model (bool, optional): If True, the model is loaded upon
            initialization. Defaults to False.
    """

    def __init__(self, settings: Settings, load_model: bool = False) -> None:
        """Initializes the MosquitoSegmenter."""
        provider_service = ProviderService(settings)
        weights_manager = ModelWeightsManager(
            settings=settings,
            provider_service=provider_service,
        )
        super().__init__(
            settings=settings,
            predictor_type="segmenter",
            weights_manager=weights_manager,
            load_model=load_model,
        )

    def predict(self, input_data: np.ndarray, **kwargs: Any) -> np.ndarray:
        """Generates a segmentation mask for mosquitos in an image.

        If `detection_boxes` are provided in kwargs, the model will generate
        masks only for those specific regions. Otherwise, it will attempt to
        segment the most prominent object in the image.

        Args:
            input_data (np.ndarray): The input image as a NumPy array (H, W, 3).
            **kwargs (Any): Optional keyword arguments, including:
                detection_boxes (list, optional): A list of bounding boxes
                in (cx, cy, w, h, conf) format to guide segmentation.

        Returns:
            np.ndarray: A binary segmentation mask of shape (H, W), where pixels
            belonging to a mosquito are marked as 1 (or True) and 0 otherwise.

        Raises:
            RuntimeError: If the model is not loaded or fails to load.
        """
        if not self.model_loaded or self._model is None:
            self.load_model()
            if self._model is None:
                raise RuntimeError("Failed to load model")

        detection_boxes = kwargs.get("detection_boxes")

        if len(input_data.shape) == 2:
            input_data = cv2.cvtColor(input_data, cv2.COLOR_GRAY2RGB)
        elif input_data.shape[2] == 4:
            input_data = cv2.cvtColor(input_data, cv2.COLOR_RGBA2RGB)

        model = cast(SAM2ImagePredictor, self._model)
        model.set_image(input_data)

        if detection_boxes and len(detection_boxes) > 0:
            masks = []
            for box in detection_boxes:
                x, y, w, h, _ = box
                x1, y1 = int(x - w / 2), int(y - h / 2)
                x2, y2 = int(x + w / 2), int(y + h / 2)
                input_box = np.array([x1, y1, x2, y2])

                mask, _, _ = model.predict(
                    point_coords=None,
                    point_labels=None,
                    box=input_box[None, :],
                    multimask_output=False,
                )
                masks.append(mask[0].astype(np.uint8))
            return np.logical_or.reduce(masks) if masks else np.zeros(input_data.shape[:2], dtype=bool)

        masks, _, _ = model.predict(point_coords=None, point_labels=None, box=None, multimask_output=False)
        return masks[0].astype(np.uint8)

    def visualize(
        self,
        input_data: np.ndarray,
        predictions: SegmentationPredictionType,
        save_path: str | Path | None = None,
    ) -> np.ndarray:
        """Overlays a segmentation mask on the original image.

        Args:
            input_data (np.ndarray): The original image.
            predictions (SegmentationPredictionType): The binary segmentation mask
                from the `predict` method.
            save_path (str | Path | None, optional): If provided, the visualized
                image is saved to this path. Defaults to None.

        Returns:
            np.ndarray: A new image array with the segmentation mask visualized
            as a colored overlay.
        """
        if len(input_data.shape) == 2:
            input_data = cv2.cvtColor(input_data, cv2.COLOR_GRAY2BGR)

        colored_mask = np.zeros_like(input_data)
        overlay_color_bgr = str_to_bgr(self.config.visualization.overlay_color)
        colored_mask[predictions > 0] = np.array(overlay_color_bgr)

        overlay = cv2.addWeighted(input_data, 1, colored_mask, self.config.visualization.alpha, 0)

        if save_path:
            cv2.imwrite(str(save_path), cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))

        return overlay

    def _evaluate_from_prediction(
        self,
        prediction: SegmentationPredictionType,
        ground_truth: SegmentationGroundTruthType,
    ) -> dict[str, float]:
        """Calculates segmentation metrics for a single predicted mask.

        Computes Intersection over Union (IoU), precision, recall, and F1-score.

        Args:
            prediction (SegmentationPredictionType): The predicted binary mask.
            ground_truth (SegmentationGroundTruthType): The ground truth binary mask.

        Returns:
            dict[str, float]: A dictionary containing the segmentation metrics.
        """
        prediction = prediction.astype(bool)
        ground_truth = ground_truth.astype(bool)

        intersection = np.logical_and(prediction, ground_truth)
        union = np.logical_or(prediction, ground_truth)

        intersection_sum = np.sum(intersection)
        prediction_sum = np.sum(prediction)
        ground_truth_sum = np.sum(ground_truth)
        union_sum = np.sum(union)

        iou = intersection_sum / union_sum if union_sum > 0 else 0.0
        precision = intersection_sum / prediction_sum if prediction_sum > 0 else 0.0
        recall = intersection_sum / ground_truth_sum if ground_truth_sum > 0 else 0.0
        f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0

        return {"iou": float(iou), "precision": float(precision), "recall": float(recall), "f1": float(f1)}

    def _load_model(self) -> None:
        """Loads the SAM2 model and initializes the image predictor.

        Raises:
            ValueError: If the required configuration for the model path or
                device is missing.
            RuntimeError: If the model fails to load from the specified path.
        """
        if (
            not hasattr(self.config, "model_config_path")
            or self.config.model_config_path is None
            or not hasattr(self.config, "device")
        ):
            raise ValueError("Missing required configuration: 'model_config_path' and 'device' must be set")

        sam2_model = build_sam2(self.config.model_config_path, str(self.model_path), device=self.config.device)
        try:
            self._model = SAM2ImagePredictor(sam2_model)
            self._model_loaded = True
        except Exception as e:
            raise RuntimeError(
                f"Failed to load SAM model from {self.model_path}. "
                f"Please check the model path and configuration. Error: {str(e)}",
            )

settings = settings instance-attribute

predictor_type = predictor_type instance-attribute

config: PredictorConfig property

Get the predictor configuration Pydantic model.

model_loaded: bool property

Check if the model is loaded.

model_path: Path property

Gets the path to the model weights file.

__call__(input_data: np.ndarray, **kwargs: Any) -> Any

Convenience method that calls predict().

Source code in culicidaelab/core/base_predictor.py

def __call__(self, input_data: np.ndarray, **kwargs: Any) -> Any:
    """Convenience method that calls `predict()`."""
    if not self._model_loaded:
        self.load_model()
    return self.predict(input_data, **kwargs)

__enter__()

Context manager entry.

Source code in culicidaelab/core/base_predictor.py

def __enter__(self):
    """Context manager entry."""
    if not self._model_loaded:
        self.load_model()
    return self

__exit__(exc_type, exc_val, exc_tb)

Context manager exit.

Source code in culicidaelab/core/base_predictor.py

def __exit__(self, exc_type, exc_val, exc_tb):
    """Context manager exit."""
    pass

model_context()

A context manager for temporary model loading.

Ensures the model is loaded upon entering the context and unloaded upon exiting. This is useful for managing memory in pipelines.

Yields:

Name	Type	Description
BasePredictor		The predictor instance itself.

Example

with predictor.model_context(): ... predictions = predictor.predict(data)

Source code in culicidaelab/core/base_predictor.py

@contextmanager
def model_context(self):
    """A context manager for temporary model loading.

    Ensures the model is loaded upon entering the context and unloaded
    upon exiting. This is useful for managing memory in pipelines.

    Yields:
        BasePredictor: The predictor instance itself.

    Example:
        >>> with predictor.model_context():
        ...     predictions = predictor.predict(data)
    """
    was_loaded = self._model_loaded
    try:
        if not was_loaded:
            self.load_model()
        yield self
    finally:
        if not was_loaded and self._model_loaded:
            self.unload_model()

evaluate(ground_truth: GroundTruthType, prediction: PredictionType | None = None, input_data: np.ndarray | None = None, **predict_kwargs: Any) -> dict[str, float]

Evaluate a prediction against a ground truth.

Either prediction or input_data must be provided. If prediction is provided, it is used directly. If prediction is None, input_data is used to generate a new prediction.

Parameters:

Name	Type	Description	Default
ground_truth	GroundTruthType	The ground truth annotation.	required
prediction	PredictionType	A pre-computed prediction.	None
input_data	ndarray	Input data to generate a prediction from, if one isn't provided.	None
**predict_kwargs	Any	Additional arguments passed to the predict method.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing evaluation metrics for a
dict[str, float]	single item.

Raises:

Type	Description
ValueError	If neither prediction nor input_data is provided.

Source code in culicidaelab/core/base_predictor.py

def evaluate(
    self,
    ground_truth: GroundTruthType,
    prediction: PredictionType | None = None,
    input_data: np.ndarray | None = None,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate a prediction against a ground truth.

    Either `prediction` or `input_data` must be provided. If `prediction`
    is provided, it is used directly. If `prediction` is None, `input_data`
    is used to generate a new prediction.

    Args:
        ground_truth (GroundTruthType): The ground truth annotation.
        prediction (PredictionType, optional): A pre-computed prediction.
        input_data (np.ndarray, optional): Input data to generate a
            prediction from, if one isn't provided.
        **predict_kwargs (Any): Additional arguments passed to the `predict`
            method.

    Returns:
        dict[str, float]: Dictionary containing evaluation metrics for a
        single item.

    Raises:
        ValueError: If neither `prediction` nor `input_data` is provided.
    """
    if prediction is None:
        if input_data is not None:
            prediction = self.predict(input_data, **predict_kwargs)
        else:
            raise ValueError(
                "Either 'prediction' or 'input_data' must be provided.",
            )
    return self._evaluate_from_prediction(
        prediction=prediction,
        ground_truth=ground_truth,
    )

evaluate_batch(ground_truth_batch: list[GroundTruthType], predictions_batch: list[PredictionType] | None = None, input_data_batch: list[np.ndarray] | None = None, num_workers: int = 4, show_progress: bool = True, **predict_kwargs: Any) -> dict[str, float]

Evaluate on a batch of items using parallel processing.

Either predictions_batch or input_data_batch must be provided.

Parameters:

Name	Type	Description	Default
ground_truth_batch	list[GroundTruthType]	List of corresponding ground truth annotations.	required
predictions_batch	list[PredictionType]	A pre-computed list of predictions.	None
input_data_batch	list[ndarray]	List of input data to generate predictions from.	None
num_workers	int	Number of parallel workers for calculating metrics.	4
show_progress	bool	Whether to show a progress bar.	True
**predict_kwargs	Any	Additional arguments passed to predict_batch.	{}

Returns:

Type	Description
dict[str, float]	dict[str, float]: Dictionary containing aggregated evaluation metrics.

Raises:

Type	Description
ValueError	If the number of predictions does not match the number of ground truths.

Source code in culicidaelab/core/base_predictor.py

def evaluate_batch(
    self,
    ground_truth_batch: list[GroundTruthType],
    predictions_batch: list[PredictionType] | None = None,
    input_data_batch: list[np.ndarray] | None = None,
    num_workers: int = 4,
    show_progress: bool = True,
    **predict_kwargs: Any,
) -> dict[str, float]:
    """Evaluate on a batch of items using parallel processing.

    Either `predictions_batch` or `input_data_batch` must be provided.

    Args:
        ground_truth_batch (list[GroundTruthType]): List of corresponding
            ground truth annotations.
        predictions_batch (list[PredictionType], optional): A pre-computed
            list of predictions.
        input_data_batch (list[np.ndarray], optional): List of input data
            to generate predictions from.
        num_workers (int): Number of parallel workers for calculating metrics.
        show_progress (bool): Whether to show a progress bar.
        **predict_kwargs (Any): Additional arguments passed to `predict_batch`.

    Returns:
        dict[str, float]: Dictionary containing aggregated evaluation metrics.

    Raises:
        ValueError: If the number of predictions does not match the number
            of ground truths.
    """
    if predictions_batch is None:
        if input_data_batch is not None:
            predictions_batch = self.predict_batch(
                input_data_batch,
                show_progress=show_progress,
                **predict_kwargs,
            )
        else:
            raise ValueError(
                "Either 'predictions_batch' or 'input_data_batch' must be provided.",
            )

    if len(predictions_batch) != len(ground_truth_batch):
        raise ValueError(
            f"Number of predictions ({len(predictions_batch)}) must match "
            f"number of ground truths ({len(ground_truth_batch)}).",
        )

    per_item_metrics = self._calculate_metrics_parallel(
        predictions_batch,
        ground_truth_batch,
        num_workers,
        show_progress,
    )
    aggregated_metrics = self._aggregate_metrics(per_item_metrics)
    final_report = self._finalize_evaluation_report(
        aggregated_metrics,
        predictions_batch,
        ground_truth_batch,
    )
    return final_report

get_model_info() -> dict[str, Any]

Gets information about the loaded model.

Returns:

Type	Description
dict[str, Any]	dict[str, Any]: A dictionary containing details about the model, such
dict[str, Any]	as architecture, path, etc.

Source code in culicidaelab/core/base_predictor.py

def get_model_info(self) -> dict[str, Any]:
    """Gets information about the loaded model.

    Returns:
        dict[str, Any]: A dictionary containing details about the model, such
        as architecture, path, etc.
    """
    return {
        "predictor_type": self.predictor_type,
        "model_path": str(self._model_path),
        "model_loaded": self._model_loaded,
        "config": self.config.model_dump(),
    }

load_model() -> None

Loads the model if it is not already loaded.

This is a convenience wrapper around _load_model that prevents reloading.

Raises:

Type	Description
RuntimeError	If model loading fails.

Source code in culicidaelab/core/base_predictor.py

def load_model(self) -> None:
    """Loads the model if it is not already loaded.

    This is a convenience wrapper around `_load_model` that prevents
    reloading.

    Raises:
        RuntimeError: If model loading fails.
    """
    if self._model_loaded:
        self._logger.info(f"Model for {self.predictor_type} already loaded")
        return

    try:
        self._logger.info(
            f"Loading model for {self.predictor_type} from {self._model_path}",
        )
        self._load_model()
        self._model_loaded = True
        self._logger.info(f"Successfully loaded model for {self.predictor_type}")
    except Exception as e:
        self._logger.error(f"Failed to load model for {self.predictor_type}: {e}")
        raise RuntimeError(
            f"Failed to load model for {self.predictor_type}: {e}",
        ) from e

predict_batch(input_data_batch: list[np.ndarray], show_progress: bool = True, **kwargs: Any) -> list[PredictionType]

Makes predictions on a batch of inputs.

This base implementation processes items serially. Subclasses with native batching capabilities SHOULD override this method.

Parameters:

Name	Type	Description	Default
input_data_batch	list[ndarray]	List of input data to make predictions on.	required
show_progress	bool	Whether to show a progress bar.	True
**kwargs	Any	Additional arguments passed to each predict call.	{}

Returns:

Type	Description
list[PredictionType]	list[PredictionType]: List of predictions.

Raises:

Type	Description
RuntimeError	If model fails to load or predict.

Source code in culicidaelab/core/base_predictor.py

def predict_batch(
    self,
    input_data_batch: list[np.ndarray],
    show_progress: bool = True,
    **kwargs: Any,
) -> list[PredictionType]:
    """Makes predictions on a batch of inputs.

    This base implementation processes items serially. Subclasses with
    native batching capabilities SHOULD override this method.

    Args:
        input_data_batch (list[np.ndarray]): List of input data to make
            predictions on.
        show_progress (bool): Whether to show a progress bar.
        **kwargs (Any): Additional arguments passed to each `predict` call.

    Returns:
        list[PredictionType]: List of predictions.

    Raises:
        RuntimeError: If model fails to load or predict.
    """
    if not input_data_batch:
        return []

    if not self._model_loaded:
        self.load_model()
        if not self._model_loaded:
            raise RuntimeError("Failed to load model for batch prediction")

    in_notebook = "ipykernel" in sys.modules
    tqdm_iterator = tqdm_notebook if in_notebook else tqdm_console
    iterator = input_data_batch

    if show_progress:
        iterator = tqdm_iterator(
            input_data_batch,
            desc=f"Predicting batch ({self.predictor_type})",
            leave=False,
        )
    try:
        return [self.predict(item, **kwargs) for item in iterator]
    except Exception as e:
        self._logger.error(f"Batch prediction failed: {e}", exc_info=True)
        raise RuntimeError(f"Batch prediction failed: {e}") from e

unload_model() -> None

Unloads the model to free memory.

Source code in culicidaelab/core/base_predictor.py

def unload_model(self) -> None:
    """Unloads the model to free memory."""
    if self._model_loaded:
        self._model = None
        self._model_loaded = False
        self._logger.info(f"Unloaded model for {self.predictor_type}")

__init__(settings: Settings, load_model: bool = False) -> None

Initializes the MosquitoSegmenter.

Source code in culicidaelab/predictors/segmenter.py

def __init__(self, settings: Settings, load_model: bool = False) -> None:
    """Initializes the MosquitoSegmenter."""
    provider_service = ProviderService(settings)
    weights_manager = ModelWeightsManager(
        settings=settings,
        provider_service=provider_service,
    )
    super().__init__(
        settings=settings,
        predictor_type="segmenter",
        weights_manager=weights_manager,
        load_model=load_model,
    )

predict(input_data: np.ndarray, **kwargs: Any) -> np.ndarray

Generates a segmentation mask for mosquitos in an image.

If detection_boxes are provided in kwargs, the model will generate masks only for those specific regions. Otherwise, it will attempt to segment the most prominent object in the image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	The input image as a NumPy array (H, W, 3).	required
**kwargs	Any	Optional keyword arguments, including: detection_boxes (list, optional): A list of bounding boxes in (cx, cy, w, h, conf) format to guide segmentation.	{}

Returns:

Type	Description
ndarray	np.ndarray: A binary segmentation mask of shape (H, W), where pixels
ndarray	belonging to a mosquito are marked as 1 (or True) and 0 otherwise.

Raises:

Type	Description
RuntimeError	If the model is not loaded or fails to load.

Source code in culicidaelab/predictors/segmenter.py

def predict(self, input_data: np.ndarray, **kwargs: Any) -> np.ndarray:
    """Generates a segmentation mask for mosquitos in an image.

    If `detection_boxes` are provided in kwargs, the model will generate
    masks only for those specific regions. Otherwise, it will attempt to
    segment the most prominent object in the image.

    Args:
        input_data (np.ndarray): The input image as a NumPy array (H, W, 3).
        **kwargs (Any): Optional keyword arguments, including:
            detection_boxes (list, optional): A list of bounding boxes
            in (cx, cy, w, h, conf) format to guide segmentation.

    Returns:
        np.ndarray: A binary segmentation mask of shape (H, W), where pixels
        belonging to a mosquito are marked as 1 (or True) and 0 otherwise.

    Raises:
        RuntimeError: If the model is not loaded or fails to load.
    """
    if not self.model_loaded or self._model is None:
        self.load_model()
        if self._model is None:
            raise RuntimeError("Failed to load model")

    detection_boxes = kwargs.get("detection_boxes")

    if len(input_data.shape) == 2:
        input_data = cv2.cvtColor(input_data, cv2.COLOR_GRAY2RGB)
    elif input_data.shape[2] == 4:
        input_data = cv2.cvtColor(input_data, cv2.COLOR_RGBA2RGB)

    model = cast(SAM2ImagePredictor, self._model)
    model.set_image(input_data)

    if detection_boxes and len(detection_boxes) > 0:
        masks = []
        for box in detection_boxes:
            x, y, w, h, _ = box
            x1, y1 = int(x - w / 2), int(y - h / 2)
            x2, y2 = int(x + w / 2), int(y + h / 2)
            input_box = np.array([x1, y1, x2, y2])

            mask, _, _ = model.predict(
                point_coords=None,
                point_labels=None,
                box=input_box[None, :],
                multimask_output=False,
            )
            masks.append(mask[0].astype(np.uint8))
        return np.logical_or.reduce(masks) if masks else np.zeros(input_data.shape[:2], dtype=bool)

    masks, _, _ = model.predict(point_coords=None, point_labels=None, box=None, multimask_output=False)
    return masks[0].astype(np.uint8)

visualize(input_data: np.ndarray, predictions: SegmentationPredictionType, save_path: str | Path | None = None) -> np.ndarray

Overlays a segmentation mask on the original image.

Parameters:

Name	Type	Description	Default
input_data	ndarray	The original image.	required
predictions	SegmentationPredictionType	The binary segmentation mask from the predict method.	required
save_path	str | Path | None	If provided, the visualized image is saved to this path. Defaults to None.	None

Returns:

Type	Description
ndarray	np.ndarray: A new image array with the segmentation mask visualized
ndarray	as a colored overlay.

Source code in culicidaelab/predictors/segmenter.py

def visualize(
    self,
    input_data: np.ndarray,
    predictions: SegmentationPredictionType,
    save_path: str | Path | None = None,
) -> np.ndarray:
    """Overlays a segmentation mask on the original image.

    Args:
        input_data (np.ndarray): The original image.
        predictions (SegmentationPredictionType): The binary segmentation mask
            from the `predict` method.
        save_path (str | Path | None, optional): If provided, the visualized
            image is saved to this path. Defaults to None.

    Returns:
        np.ndarray: A new image array with the segmentation mask visualized
        as a colored overlay.
    """
    if len(input_data.shape) == 2:
        input_data = cv2.cvtColor(input_data, cv2.COLOR_GRAY2BGR)

    colored_mask = np.zeros_like(input_data)
    overlay_color_bgr = str_to_bgr(self.config.visualization.overlay_color)
    colored_mask[predictions > 0] = np.array(overlay_color_bgr)

    overlay = cv2.addWeighted(input_data, 1, colored_mask, self.config.visualization.alpha, 0)

    if save_path:
        cv2.imwrite(str(save_path), cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))

    return overlay

ModelWeightsManager

Manages the download and local availability of model weights.

This class implements the WeightsManagerProtocol and serves as the bridge between a predictor and the provider service that can download model files.

Parameters:

Name	Type	Description	Default
settings	Settings	The application's global settings object.	required
provider_service	ProviderService	The service for downloading artifacts.	required

Source code in culicidaelab/predictors/model_weights_manager.py

class ModelWeightsManager(WeightsManagerProtocol):
    """Manages the download and local availability of model weights.

    This class implements the WeightsManagerProtocol and serves as the bridge
    between a predictor and the provider service that can download model files.

    Args:
        settings (Settings): The application's global settings object.
        provider_service (ProviderService): The service for downloading artifacts.
    """

    def __init__(self, settings: Settings, provider_service: ProviderService):
        """Initializes the ModelWeightsManager."""
        self.settings = settings
        self.provider_service = provider_service

    def ensure_weights(self, model_type: str) -> Path:
        """Ensures model weights exist locally, downloading them if needed.

        This method checks for the local existence of a model's weights. If they
        are not found, it uses the provider service to download them based on
        the configuration associated with the given model type.

        Args:
            model_type (str): The type of the model for which to ensure weights,
                e.g., 'classifier', 'detector'.

        Returns:
            Path: The absolute, canonical path to the local model weights file.

        Raises:
            RuntimeError: If the weights cannot be downloaded or if the
                configuration for the provider is missing or invalid.
        """
        predictor_config = None
        try:
            predictor_config = self.settings.get_config(f"predictors.{model_type}")
            provider = self.provider_service.get_provider(predictor_config.provider_name)
            return provider.download_model_weights(model_type)
        except Exception as e:
            error_msg = f"Failed to download weights for '{model_type}': {str(e)}"
            if predictor_config:
                error_msg += f" with predictor config {predictor_config}"
            raise RuntimeError(error_msg) from e

settings = settings instance-attribute

provider_service = provider_service instance-attribute

__init__(settings: Settings, provider_service: ProviderService)

Initializes the ModelWeightsManager.

Source code in culicidaelab/predictors/model_weights_manager.py

def __init__(self, settings: Settings, provider_service: ProviderService):
    """Initializes the ModelWeightsManager."""
    self.settings = settings
    self.provider_service = provider_service

ensure_weights(model_type: str) -> Path

Ensures model weights exist locally, downloading them if needed.

This method checks for the local existence of a model's weights. If they are not found, it uses the provider service to download them based on the configuration associated with the given model type.

Parameters:

Name	Type	Description	Default
model_type	str	The type of the model for which to ensure weights, e.g., 'classifier', 'detector'.	required

Returns:

Name	Type	Description
Path	Path	The absolute, canonical path to the local model weights file.

Raises:

Type	Description
RuntimeError	If the weights cannot be downloaded or if the configuration for the provider is missing or invalid.

Source code in culicidaelab/predictors/model_weights_manager.py

def ensure_weights(self, model_type: str) -> Path:
    """Ensures model weights exist locally, downloading them if needed.

    This method checks for the local existence of a model's weights. If they
    are not found, it uses the provider service to download them based on
    the configuration associated with the given model type.

    Args:
        model_type (str): The type of the model for which to ensure weights,
            e.g., 'classifier', 'detector'.

    Returns:
        Path: The absolute, canonical path to the local model weights file.

    Raises:
        RuntimeError: If the weights cannot be downloaded or if the
            configuration for the provider is missing or invalid.
    """
    predictor_config = None
    try:
        predictor_config = self.settings.get_config(f"predictors.{model_type}")
        provider = self.provider_service.get_provider(predictor_config.provider_name)
        return provider.download_model_weights(model_type)
    except Exception as e:
        error_msg = f"Failed to download weights for '{model_type}': {str(e)}"
        if predictor_config:
            error_msg += f" with predictor config {predictor_config}"
        raise RuntimeError(error_msg) from e

selection:

members: true