ibbi 0.0.0__tar.gz

ibbi-0.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Chris
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ibbi-0.0.0/PKG-INFO

@@ -0,0 +1,200 @@
+Metadata-Version: 2.3
+Name: ibbi
+Version: 0.0.0
+Summary: A package for bark and ambrosia beetle identification.
+License: MIT
+Author: G. Christopher Marais
+Requires-Python: >=3.11,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: huggingface-hub (>=0.32.3,<0.33.0)
+Requires-Dist: ipywidgets (>=8.1.7,<9.0.0)
+Requires-Dist: numpy (>=2.2.6,<3.0.0)
+Requires-Dist: pandas (>=2.3.0,<3.0.0)
+Requires-Dist: pillow (>=11.2.1,<12.0.0)
+Requires-Dist: ultralytics (==8.3.139)
+Description-Content-Type: text/markdown
+
+# Intelligent Bark Beetle Identifier (IBBI)
+
+<!-- [![JOSS submission](https://joss.theoj.org/papers/10.21105/joss.01234/status.svg)](https://joss.theoj.org/papers/10.21105/joss.01234) -->
+[![PyPI version](https://badge.fury.io/py/ibbi.svg)](https://badge.fury.io/py/ibbi)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+**IBBI** is a Python package that provides a simple and unified interface for detecting and classifying bark and ambrosia beetles from images using state-of-the-art computer vision models.
+
+This package is designed to support entomological research by automating the laborious task of beetle identification, enabling high-throughput data analysis for ecological studies, pest management, and biodiversity monitoring. The core models are built on multiple different architectures and are made easily accessible through a simple Python API.
+
+### Motivation
+
+The ability to accurately identify bark and ambrosia beetles is critical for forest health and pest management. However, traditional methods face significant challenges:
+
+* **They are slow and time-consuming.**
+* **They require highly specialized expertise.**
+* **They create a bottleneck for large-scale research.**
+
+The IBBI package provides a powerful, modern solution to overcome these obstacles:
+
+* It uses **pre-trained, open-source models** for rapid analysis.
+* It **automates both detection and classification** from images.
+* It **lowers the barrier to entry**, enabling faster and more extensive data collection for all researchers.
+
+---
+
+## Table of Contents
+
+- [Intelligent Bark Beetle Identifier (IBBI)](#intelligent-bark-beetle-identifier-ibbi)
+    - [Motivation](#motivation)
+  - [Table of Contents](#table-of-contents)
+  - [Installation](#installation)
+  - [Quick Start](#quick-start)
+  - [Available Models](#available-models)
+  - [Model Training Workflow](#model-training-workflow)
+  - [How to Contribute](#how-to-contribute)
+  - [License](#license)
+
+---
+
+## Installation
+
+This package requires PyTorch. For compatibility with your specific hardware (e.g., CUDA-enabled GPU), please install PyTorch *before* installing `ibbi`.
+
+**1. Install PyTorch**
+
+Follow the official instructions at **[pytorch.org](https://pytorch.org/get-started/locally/)** to install the correct version for your system (OS, package manager, and CUDA version).
+
+**2. Install IBBI**
+
+Once PyTorch is installed, you can install the package directly from PyPI:
+
+```bash
+pip install ibbi
+````
+
+-----
+
+## Quick Start
+
+Using IBBI is straightforward. You can load a pre-trained model for either detection or classification and immediately use it for inference on your images.
+
+```python
+import ibbi
+from PIL import Image
+
+# Load an image
+image = Image.open("path/to/your/beetle_image.jpg")
+
+# 1. Load a pretrained object detection model or a classification model
+detector = ibbi.create_model("yolov10x_bb_detect_model", pretrained=True)
+classifier = ibbi.create_model("yolov10x_bb_classify_model", pretrained=True)
+
+# 2. Run prediction to get class probabilities and/or bounding boxes
+# The results will be the detected bounding box coordinates, confidence scores, and class labels
+detection_results = detector.predict(image)
+classification_results = classifier.predict(image)
+
+# 3. You can also extract deep features from all models for other tasks
+# The results will be a tensor of features
+features = classifier.extract_features(image)
+
+```
+
+For a more detailed, hands-on demonstration, please see the example notebook located in the repository: `notebooks/example.ipynb`.
+
+-----
+
+## Available Models
+
+The package provides a factory function `create_model()` to access the following pre-trained models from Huggingface Hub:
+
+| Model Name                 | Task             | Pretrained Weights Repository              | Model Size (Params) | mAP@0.5 | mAP@[.5:.95] |
+|----------------------------|------------------|--------------------------------------------|---------------------|---------|--------------|
+| yolov10x_bb_detect_model   | Object Detection | ChristopherMarais/ibbi_yolov10_od_20250601 | 29.5M                | N/A     | N/A          |
+| yolov10x_bb_classify_model | Classification   | ChristopherMarais/ibbi_yolov10_c_20250608  | 29.5M                 | N/A     | N/A          |
+
+A detailed list of available models and their Hugging Face repositories can be found in the [ibbi_model_summary.csv](./docs/assets/data/ibbi_model_summary.csv) file.
+
+-----
+
+## Model Training Workflow
+
+The models included in this package were trained using a standardized data flow on a dataset of bark and ambrosia beetle images from multiple different sources to include 63 different species. A list of all the species the current set of classification models were trained on can be found in the [ibbi_species_table.csv](./docs/assets/data/ibbi_species_table.csv) file.
+
+<p align="center">
+  <img src="./docs/assets/images/data_flow_ibbi.png" alt="My training workflow">
+</p>
+
+**1. Data Collection & Aggregation:**
+
+  * An initial dataset of 54,421 images was compiled from diverse sources, including field photography from [barkbeetles.info](https://www.barkbeetles.info), lab-based specimen photography, and images from iNaturalist.
+  * This aggregated dataset contained a mix of labeled and unlabeled images. Initially, 17,689 images had species-level labels, while 36,732 had no annotations.
+
+**2. Annotation with Human-in-the-Loop:**
+
+  * To create high-quality localization data, a zero-shot object detection model (GroundingDINO) was first used to generate preliminary bounding boxes for the beetles in the images.
+  * Crucially, these automated annotations were then manually reviewed and refined by experts to ensure their accuracy and consistency, creating a reliable ground truth for training.
+
+**3. Dataset Preparation and Splitting:**
+
+  * A dedicated test set of 2,031 images was created by selecting images only from species with at least 50 representatives. This ensures a balanced and fair evaluation.
+  * The remaining annotated data was split into two distinct training sets based on the task:
+    * Object Detection Training Set (35,274 images): All images with verified bounding boxes (excluding the test set) were used to train the general beetle detection models. This larger dataset helps the models learn to accurately localize beetles under various conditions.
+    * Classification Training Set (11,507 images): A filtered subset containing images with both verified bounding boxes and species-level labels was used to train the fine-grained classification models.
+
+**4. Model Training and Fine-Tuning:**
+
+  * Pre-trained model architectures were fine-tuned for each specific task:
+    * Object Detection models were trained on the larger localization dataset to become expert beetle detectors.
+    * Classification models were trained on the fully labeled dataset to specialize in identifying different beetle species.
+  * Data augmentation techniques, including random rotations, scaling, color jitter, and mosaic augmentation, were used throughout training to improve model robustness and prevent over-fitting.
+
+**5. Evaluation and Deployment:**
+
+  * The performance of all trained models was rigorously measured against the held-out test set to ensure high accuracy for both detection and classification tasks.
+  * The final, best-performing model weights were saved and uploaded to Hugging Face Hub, from where the `ibbi` package automatically downloads them, making them easily accessible to researchers via a simple API.
+
+-----
+
+## How to Contribute
+
+Contributions are welcome\! If you would like to improve IBBI, please follow these steps:
+
+1.  Clone this repository.
+2.  Create a Conda environment and activate it:
+    ```bash
+    conda env create -f environment.yml
+    conda activate IBBI
+    ```
+3.  Install dependencies using Poetry and set up pre-commit hooks:
+    ```bash
+    pip install torch torchvision torchaudio # Ensure PyTorch is installed first
+    poetry config virtualenvs.create false --local
+    poetry install
+    poetry run pre-commit install
+    ```
+4.  Create a new branch for your feature or bug fix.
+5.  Commit your changes and open a pull request.
+
+To add new dependencies, use `poetry add <package-name>` for the main package or `poetry add --group dev <package-name>` for development dependencies.
+
+-----
+
+<!-- ## Citing IBBI
+
+If you use IBBI in your research, please cite the JOSS paper.
+
+**(Placeholder) To be added upon acceptance:**
+
+> Marais, C., et al., (2025). IBBI: Intelligent Bark Beetle Identifier. Journal of Open Source Software, X(XX), XXXX. https://www.google.com/search?q=https://doi.org/XX.XXXXX/joss.XXXXX
+
+You can also cite the specific version of the software archive using the DOI provided by Zenodo/figshare.
+
+----- -->
+
+## License
+
+This project is licensed under the terms of the MIT License. See the `LICENSE` file for details.
+

ibbi-0.0.0/README.md

@@ -0,0 +1,179 @@
+# Intelligent Bark Beetle Identifier (IBBI)
+
+<!-- [![JOSS submission](https://joss.theoj.org/papers/10.21105/joss.01234/status.svg)](https://joss.theoj.org/papers/10.21105/joss.01234) -->
+[![PyPI version](https://badge.fury.io/py/ibbi.svg)](https://badge.fury.io/py/ibbi)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+**IBBI** is a Python package that provides a simple and unified interface for detecting and classifying bark and ambrosia beetles from images using state-of-the-art computer vision models.
+
+This package is designed to support entomological research by automating the laborious task of beetle identification, enabling high-throughput data analysis for ecological studies, pest management, and biodiversity monitoring. The core models are built on multiple different architectures and are made easily accessible through a simple Python API.
+
+### Motivation
+
+The ability to accurately identify bark and ambrosia beetles is critical for forest health and pest management. However, traditional methods face significant challenges:
+
+* **They are slow and time-consuming.**
+* **They require highly specialized expertise.**
+* **They create a bottleneck for large-scale research.**
+
+The IBBI package provides a powerful, modern solution to overcome these obstacles:
+
+* It uses **pre-trained, open-source models** for rapid analysis.
+* It **automates both detection and classification** from images.
+* It **lowers the barrier to entry**, enabling faster and more extensive data collection for all researchers.
+
+---
+
+## Table of Contents
+
+- [Intelligent Bark Beetle Identifier (IBBI)](#intelligent-bark-beetle-identifier-ibbi)
+    - [Motivation](#motivation)
+  - [Table of Contents](#table-of-contents)
+  - [Installation](#installation)
+  - [Quick Start](#quick-start)
+  - [Available Models](#available-models)
+  - [Model Training Workflow](#model-training-workflow)
+  - [How to Contribute](#how-to-contribute)
+  - [License](#license)
+
+---
+
+## Installation
+
+This package requires PyTorch. For compatibility with your specific hardware (e.g., CUDA-enabled GPU), please install PyTorch *before* installing `ibbi`.
+
+**1. Install PyTorch**
+
+Follow the official instructions at **[pytorch.org](https://pytorch.org/get-started/locally/)** to install the correct version for your system (OS, package manager, and CUDA version).
+
+**2. Install IBBI**
+
+Once PyTorch is installed, you can install the package directly from PyPI:
+
+```bash
+pip install ibbi
+````
+
+-----
+
+## Quick Start
+
+Using IBBI is straightforward. You can load a pre-trained model for either detection or classification and immediately use it for inference on your images.
+
+```python
+import ibbi
+from PIL import Image
+
+# Load an image
+image = Image.open("path/to/your/beetle_image.jpg")
+
+# 1. Load a pretrained object detection model or a classification model
+detector = ibbi.create_model("yolov10x_bb_detect_model", pretrained=True)
+classifier = ibbi.create_model("yolov10x_bb_classify_model", pretrained=True)
+
+# 2. Run prediction to get class probabilities and/or bounding boxes
+# The results will be the detected bounding box coordinates, confidence scores, and class labels
+detection_results = detector.predict(image)
+classification_results = classifier.predict(image)
+
+# 3. You can also extract deep features from all models for other tasks
+# The results will be a tensor of features
+features = classifier.extract_features(image)
+
+```
+
+For a more detailed, hands-on demonstration, please see the example notebook located in the repository: `notebooks/example.ipynb`.
+
+-----
+
+## Available Models
+
+The package provides a factory function `create_model()` to access the following pre-trained models from Huggingface Hub:
+
+| Model Name                 | Task             | Pretrained Weights Repository              | Model Size (Params) | mAP@0.5 | mAP@[.5:.95] |
+|----------------------------|------------------|--------------------------------------------|---------------------|---------|--------------|
+| yolov10x_bb_detect_model   | Object Detection | ChristopherMarais/ibbi_yolov10_od_20250601 | 29.5M                | N/A     | N/A          |
+| yolov10x_bb_classify_model | Classification   | ChristopherMarais/ibbi_yolov10_c_20250608  | 29.5M                 | N/A     | N/A          |
+
+A detailed list of available models and their Hugging Face repositories can be found in the [ibbi_model_summary.csv](./docs/assets/data/ibbi_model_summary.csv) file.
+
+-----
+
+## Model Training Workflow
+
+The models included in this package were trained using a standardized data flow on a dataset of bark and ambrosia beetle images from multiple different sources to include 63 different species. A list of all the species the current set of classification models were trained on can be found in the [ibbi_species_table.csv](./docs/assets/data/ibbi_species_table.csv) file.
+
+<p align="center">
+  <img src="./docs/assets/images/data_flow_ibbi.png" alt="My training workflow">
+</p>
+
+**1. Data Collection & Aggregation:**
+
+  * An initial dataset of 54,421 images was compiled from diverse sources, including field photography from [barkbeetles.info](https://www.barkbeetles.info), lab-based specimen photography, and images from iNaturalist.
+  * This aggregated dataset contained a mix of labeled and unlabeled images. Initially, 17,689 images had species-level labels, while 36,732 had no annotations.
+
+**2. Annotation with Human-in-the-Loop:**
+
+  * To create high-quality localization data, a zero-shot object detection model (GroundingDINO) was first used to generate preliminary bounding boxes for the beetles in the images.
+  * Crucially, these automated annotations were then manually reviewed and refined by experts to ensure their accuracy and consistency, creating a reliable ground truth for training.
+
+**3. Dataset Preparation and Splitting:**
+
+  * A dedicated test set of 2,031 images was created by selecting images only from species with at least 50 representatives. This ensures a balanced and fair evaluation.
+  * The remaining annotated data was split into two distinct training sets based on the task:
+    * Object Detection Training Set (35,274 images): All images with verified bounding boxes (excluding the test set) were used to train the general beetle detection models. This larger dataset helps the models learn to accurately localize beetles under various conditions.
+    * Classification Training Set (11,507 images): A filtered subset containing images with both verified bounding boxes and species-level labels was used to train the fine-grained classification models.
+
+**4. Model Training and Fine-Tuning:**
+
+  * Pre-trained model architectures were fine-tuned for each specific task:
+    * Object Detection models were trained on the larger localization dataset to become expert beetle detectors.
+    * Classification models were trained on the fully labeled dataset to specialize in identifying different beetle species.
+  * Data augmentation techniques, including random rotations, scaling, color jitter, and mosaic augmentation, were used throughout training to improve model robustness and prevent over-fitting.
+
+**5. Evaluation and Deployment:**
+
+  * The performance of all trained models was rigorously measured against the held-out test set to ensure high accuracy for both detection and classification tasks.
+  * The final, best-performing model weights were saved and uploaded to Hugging Face Hub, from where the `ibbi` package automatically downloads them, making them easily accessible to researchers via a simple API.
+
+-----
+
+## How to Contribute
+
+Contributions are welcome\! If you would like to improve IBBI, please follow these steps:
+
+1.  Clone this repository.
+2.  Create a Conda environment and activate it:
+    ```bash
+    conda env create -f environment.yml
+    conda activate IBBI
+    ```
+3.  Install dependencies using Poetry and set up pre-commit hooks:
+    ```bash
+    pip install torch torchvision torchaudio # Ensure PyTorch is installed first
+    poetry config virtualenvs.create false --local
+    poetry install
+    poetry run pre-commit install
+    ```
+4.  Create a new branch for your feature or bug fix.
+5.  Commit your changes and open a pull request.
+
+To add new dependencies, use `poetry add <package-name>` for the main package or `poetry add --group dev <package-name>` for development dependencies.
+
+-----
+
+<!-- ## Citing IBBI
+
+If you use IBBI in your research, please cite the JOSS paper.
+
+**(Placeholder) To be added upon acceptance:**
+
+> Marais, C., et al., (2025). IBBI: Intelligent Bark Beetle Identifier. Journal of Open Source Software, X(XX), XXXX. https://www.google.com/search?q=https://doi.org/XX.XXXXX/joss.XXXXX
+
+You can also cite the specific version of the software archive using the DOI provided by Zenodo/figshare.
+
+----- -->
+
+## License
+
+This project is licensed under the terms of the MIT License. See the `LICENSE` file for details.

ibbi-0.0.0/pyproject.toml

@@ -0,0 +1,84 @@
+[tool.poetry]
+name = "ibbi"
+version = "0.0.0"
+description = "A package for bark and ambrosia beetle identification."
+authors = ["G. Christopher Marais"]
+license = "MIT"
+readme = "README.md"
+packages = [{ include = "ibbi", from = "src" }]
+
+[tool.poetry.dependencies]
+python = ">=3.11,<4.0"
+huggingface-hub = ">=0.32.3,<0.33.0"
+ultralytics = "8.3.139"
+pandas = "^2.3.0"
+numpy = "^2.2.6"
+pillow = "^11.2.1"
+ipywidgets = "^8.1.7"
+
+[tool.poetry.group.dev.dependencies]
+ruff = "^0.11.10"
+pyright = "^1.1.401"
+pytest = "^8.3.5"
+pre-commit = "^4.2.0"
+mkdocs = "^1.6.1"
+mkdocs-material = "^9.6.14"
+jupyterlab = "^4.4.2"
+isort = "^6.0.1"
+black = "^25.1.0"
+notebook = "^7.4.2"
+ipykernel = "^6.29.5"
+ipywidgets = "^8.1.7"
+mkdocstrings = "^0.29.1"
+mkdocstrings-python = "^1.16.12"
+poetry-dynamic-versioning = "^1.8.2"
+
+[build-system]
+requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning"]
+build-backend = "poetry.core.masonry.api"
+
+
+# --- Tool Configurations ---
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+select = ["A", "B", "C4", "E", "W", "F", "I", "UP", "RUF"]
+ignore = [
+    "B007",
+    "E721",
+    "RUF015",
+    "UP038",
+    "C401",
+    "C408",
+    "C409",
+]
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.ruff.lint.isort]
+known-first-party = ["ibbi"]
+
+
+[tool.pyright]
+include = ["src/ibbi", "tests"]
+exclude = ["**/node_modules", "**/__pycache__"]
+stubPath = "src/stubs"
+reportMissingImports = true
+pythonVersion = "3.11"
+
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+addopts = "-ra -q"
+testpaths = [
+    "tests",
+]
+
+[tool.poetry-dynamic-versioning]
+enable = true
+vcs = "git"
+style = "pep440"
+format-latest = "{base}.dev{distance}"

ibbi-0.0.0/src/ibbi/__init__.py

@@ -0,0 +1,68 @@
+# src/ibbi/__init__.py
+
+"""
+Main initialization file for the ibbi package.
+
+This file exposes the primary user-facing function, `create_model`, which acts as a
+factory for instantiating various beetle detection and classification models.
+"""
+
+from typing import Any, Union
+
+# --- IMPORTANT ---
+# Import the model definition files first.
+# This ensures that the @register_model decorator runs and populates
+# the model_registry before we try to use it.
+from .models import classification, detection  # noqa: F401
+
+# Now, import the registry that has been populated.
+from .models._registry import model_registry
+from .models.classification import YOLOv10BeetleClassifier
+from .models.detection import YOLOv10BeetleDetector
+
+# Define a type hint for the models that can be returned
+ModelType = Union[YOLOv10BeetleDetector, YOLOv10BeetleClassifier]
+
+
+def create_model(model_name: str, pretrained: bool = False, **kwargs: Any) -> ModelType:
+    """
+    Creates a model from a name.
+
+    This factory function is the main entry point for users of the package.
+    It looks up the requested model in the registry, downloads pretrained
+    weights from the Hugging Face Hub if requested, and returns an
+    instantiated model object.
+
+    Args:
+        model_name (str): Name of the model to create.
+        pretrained (bool): Whether to load pretrained weights from the Hugging Face Hub.
+                           Defaults to False.
+        **kwargs (Any): Extra arguments to pass to the model-creating function.
+
+    Returns:
+        ModelType: An instance of the requested model (e.g., YOLOv10BeetleDetector or
+                   YOLOv10BeetleClassifier).
+
+    Raises:
+        KeyError: If the requested `model_name` is not found in the model registry.
+
+    Example:
+        ```python
+        import ibbi
+
+        # Create a pretrained detection model
+        detector = ibbi.create_model("yolov10x_bb_detect_model", pretrained=True)
+
+        # Create a pretrained classification model
+        classifier = ibbi.create_model("yolov10x_bb_classify_model", pretrained=True)
+        ```
+    """
+    if model_name not in model_registry:
+        available = ", ".join(model_registry.keys())
+        raise KeyError(f"Model '{model_name}' not found. Available models: [{available}]")
+
+    # Look up the factory function in the registry and call it
+    model_factory = model_registry[model_name]
+    model = model_factory(pretrained=pretrained, **kwargs)
+
+    return model

ibbi-0.0.0/src/ibbi/models/__init__.py

@@ -0,0 +1,6 @@
+# src/ibbi/models/__init__.py
+
+from .classification import yolov10x_bb_classify_model
+from .detection import yolov10x_bb_detect_model
+
+__all__ = ["yolov10x_bb_detect_model", "yolov10x_bb_classify_model"]

ibbi-0.0.0/src/ibbi/models/_registry.py

@@ -0,0 +1,24 @@
+# 1. The registry itself: a dictionary to hold your models.
+model_registry = {}
+
+
+def register_model(fn):
+    """
+    # 2. A decorator function to easily add models to the registry.
+
+    This function takes a model-creating function (like your
+    `yolov10x_bb_detect_model` function) and adds it to the
+    `model_registry` dictionary. The function's name becomes the key.
+
+    Args:
+        fn: The model-creating function to register.
+
+    Returns:
+        The original function, after it has been registered.
+    """
+    model_name = fn.__name__
+    if model_name in model_registry:
+        raise ValueError(f"Model {model_name} is already registered.")
+
+    model_registry[model_name] = fn
+    return fn

ibbi-0.0.0/src/ibbi/models/classification.py

@@ -0,0 +1,102 @@
+# src/ibbi/models/classification.py
+
+"""
+Beetle classification models.
+"""
+
+import torch
+from ultralytics import YOLO
+
+from ..utils.hub import download_from_hf_hub
+from ._registry import register_model
+
+
+class YOLOv10BeetleClassifier:
+    """
+    A wrapper class for the YOLOv10 beetle classifier model.
+
+    Provides a clean interface for image classification and feature extraction.
+
+    Attributes:
+        model (YOLO): The underlying `ultralytics.YOLO` model instance.
+        device (str): The compute device ('cuda' or 'cpu') the model is loaded on.
+    """
+
+    def __init__(self, model_path: str):
+        """
+        Initializes the YOLOv10BeetleClassifier.
+
+        Args:
+            model_path (str): The local path to the .pt model file.
+        """
+        self.model = YOLO(model_path)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        print(f"Model loaded on device: {self.device}")
+
+    def predict(self, image, **kwargs):
+        """
+        Performs image classification on an image.
+
+        This method takes an image and returns the predicted class probabilities.
+        Accepts any arguments that the `ultralytics.YOLO.predict` method accepts.
+
+        Args:
+            image: The input image(s). Can be a path, URL, numpy array, PIL image, etc.
+            **kwargs: Additional keyword arguments to pass to the underlying
+                      `predict` method of the YOLO model.
+
+        Returns:
+            A list of `ultralytics.engine.results.Results` objects containing the
+            predicted class probabilities.
+        """
+        print("Running image classification (predict)...")
+        return self.model.predict(image, **kwargs)
+
+    def extract_features(self, image, **kwargs):
+        """
+        Extracts deep features from the backbone of the model for an image.
+
+        This is useful for downstream tasks like clustering or similarity search.
+
+        Args:
+            image: The input image(s).
+            **kwargs: Additional keyword arguments to pass to the underlying
+                      `embed` method of the YOLO model.
+
+        Returns:
+            A tensor of features if successful, otherwise None.
+        """
+        print("Extracting features (embed)...")
+        features = self.model.embed(image, **kwargs)
+        if features:
+            return features[0]
+        return None
+
+
+@register_model
+def yolov10x_bb_classify_model(pretrained: bool = False, **kwargs):
+    """
+    Factory function for the YOLOv10 beetle classifier.
+
+    Instantiates a `YOLOv10BeetleClassifier` model. If `pretrained` is True,
+    it downloads the official weights from the Hugging Face Hub.
+
+    Args:
+        pretrained (bool): If True, downloads pretrained weights.
+                           Defaults to False.
+        **kwargs: Additional arguments (not currently used).
+
+    Returns:
+        YOLOv10BeetleClassifier: An instance of the classifier class.
+    """
+    if pretrained:
+        repo_id = "ChristopherMarais/ibbi_yolov10_c_20250608"
+        filename = "model.pt"
+        local_weights_path = download_from_hf_hub(repo_id=repo_id, filename=filename)
+    else:
+        # Note: A non-pretrained YOLOv10x model is not meaningful for this task,
+        # but this path is kept for API consistency.
+        local_weights_path = "yolov10x.pt"
+
+    return YOLOv10BeetleClassifier(model_path=local_weights_path)

ibbi-0.0.0/src/ibbi/models/detection.py

@@ -0,0 +1,103 @@
+# src/ibbi/models/detection.py
+
+"""
+Beetle detection models.
+"""
+
+import torch
+from ultralytics import YOLO
+
+from ..utils.hub import download_from_hf_hub
+from ._registry import register_model
+
+
+class YOLOv10BeetleDetector:
+    """
+    A wrapper class for YOLOv10 beetle detection models.
+
+    This class provides a clean interface for performing object detection inference
+    and extracting deep features from images.
+
+    Attributes:
+        model (YOLO): The underlying `ultralytics.YOLO` model instance.
+        device (str): The compute device ('cuda' or 'cpu') the model is loaded on.
+    """
+
+    def __init__(self, model_path: str):
+        """
+        Initializes the YOLOv10BeetleDetector.
+
+        Args:
+            model_path (str): The local path to the pretrained `.pt` model file.
+        """
+        self.model = YOLO(model_path)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        print(f"Model loaded on device: {self.device}")
+
+    def predict(self, image, **kwargs):
+        """
+        Performs object detection inference on an image.
+
+        This method takes an image and returns the bounding box predictions.
+        Accepts any arguments that the `ultralytics.YOLO.predict` method accepts.
+
+        Args:
+            image: The input image(s). Can be a path, URL, numpy array, PIL image, etc.
+            **kwargs: Additional keyword arguments to pass to the underlying
+                      `predict` method of the YOLO model.
+
+        Returns:
+            A list of `ultralytics.engine.results.Results` objects containing the
+            detected bounding boxes, confidence scores, and class labels.
+        """
+        print("Running object detection (predict)...")
+        return self.model.predict(image, **kwargs)
+
+    def extract_features(self, image, **kwargs):
+        """
+        Extracts deep features from the backbone of the model for an image.
+
+        This is useful for downstream tasks like clustering or similarity search.
+
+        Args:
+            image: The input image(s).
+            **kwargs: Additional keyword arguments to pass to the underlying
+                      `embed` method of the YOLO model.
+
+        Returns:
+            A tensor of features if successful, otherwise None.
+        """
+        print("Extracting features (embed)...")
+        features = self.model.embed(image, **kwargs)
+        if features:
+            return features[0]
+        return None
+
+
+@register_model
+def yolov10x_bb_detect_model(pretrained: bool = False, **kwargs):
+    """
+    Factory function for the YOLOv10 beetle detector.
+
+    Instantiates a `YOLOv10BeetleDetector` model. If `pretrained` is True,
+    it downloads the official weights from the Hugging Face Hub.
+
+    Args:
+        pretrained (bool): If True, downloads pretrained weights.
+                           Defaults to False.
+        **kwargs: Additional arguments (not currently used).
+
+    Returns:
+        YOLOv10BeetleDetector: An instance of the detector class.
+    """
+    if pretrained:
+        repo_id = "ChristopherMarais/ibbi_yolov10_od_20250601"
+        filename = "model.pt"
+        local_weights_path = download_from_hf_hub(repo_id=repo_id, filename=filename)
+    else:
+        # Note: A non-pretrained YOLOv10x model is not meaningful for this task,
+        # but this path is kept for API consistency.
+        local_weights_path = "yolov10x.pt"
+
+    return YOLOv10BeetleDetector(model_path=local_weights_path)

ibbi-0.0.0/src/ibbi/utils/hub.py

@@ -0,0 +1,18 @@
+from huggingface_hub import hf_hub_download
+
+
+def download_from_hf_hub(repo_id: str, filename: str) -> str:
+    """
+    Downloads a model file from a Hugging Face Hub repository.
+
+    Args:
+        repo_id (str): The ID of the repository (e.g., "your-username/my-model").
+        filename (str): The name of the file to download from the repo.
+
+    Returns:
+        str: The local path to the downloaded file.
+    """
+    print(f"Downloading {filename} from Hugging Face hub repository '{repo_id}'...")
+    local_model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    print("Download complete. Model cached at:", local_model_path)
+    return local_model_path