README.md

---
license: mit
language:
- en
library_name: transformers
base_model:
- facebook/mask2former-swin-large-ade-semantic
tags:
- biology
- CV
- images
- animals
- beetles
- image-segmentation
- mask2former
- entomology
datasets:
- imageomics/sentinel-beetles
metrics:
- mean_iou
model_description: >-
  Fine-tuned Mask2Former models for semantic segmentation of beetle body parts
  on pinned specimen images. Two checkpoints are provided: a 5-class model
  (head, pronotum, elytra, legs, antennas) and a 9-class model (head, eyes,
  mouthparts, pronotum, elytra, tail, legs, antennas, pin). Part of the
  BeetleFlow pipeline for large-scale beetle image processing.
---

# Model Card for BeetleFlow

BeetleFlow provides fine-tuned **Mask2Former** checkpoints for **semantic segmentation of beetle body parts** on cropped, pinned specimen images. This repository hosts two variants:

| Checkpoint | Segmented parts (excluding background) | Labels |
|------------|------------------------------------------|--------|
| [`5-class`](./5-class) | 5 parts | background, head, pronotum, elytra, legs, antennas |
| [`9-class`](./9-class) | 9 parts | background, head, eyes, mouthparts, pronotum, elytra, tail, legs, antennas, pin |

These models are the segmentation stage of [BeetleFlow](https://github.com/Imageomics/BeetleFlow): an integrative deep learning pipeline for beetle image processing (detection → cropping → segmentation).

## Model Details

### Model Description

Mask2Former with a Swin-Large backbone, fine-tuned for pixel-level semantic segmentation of morphological regions on pinned ground-beetle images. Models are initialized from [`facebook/mask2former-swin-large-ade-semantic`](https://huggingface.co/facebook/mask2former-swin-large-ade-semantic) and trained on manually annotated RGB masks derived from the [sentinel-beetles](https://huggingface.co/datasets/imageomics/sentinel-beetles) dataset.

- **Developed by:** Fangxun Liu
- **Model type:** Semantic segmentation (Mask2Former)
- **Language(s) (NLP):** N/A (computer vision model)
- **License:** [MIT](https://opensource.org/licenses/MIT) (model weights and code); training images are [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/)
- **Fine-tuned from model:** [`facebook/mask2former-swin-large-ade-semantic`](https://huggingface.co/facebook/mask2former-swin-large-ade-semantic)

### Model Sources

- **Repository:** [Imageomics/BeetleFlow](https://github.com/Imageomics/BeetleFlow)
- **Paper:** [BeetleFlow: An Integrative Deep Learning Pipeline for Beetle Image Processing](https://doi.org/10.48550/arXiv.2511.00255) — *NeurIPS 2025 Workshop for Imageomics: Discovering Biological Knowledge from Images Using AI*
- **Demo:** See [batch inference usage](#how-to-get-started-with-the-model) below

## Uses

### Direct Use

- Segment beetle body parts on cropped, single-specimen images produced by the BeetleFlow detection and cropping pipeline.
- Generate per-pixel part masks for downstream morphological measurements and ecological analyses.
- Run batch inference over folders of cropped beetle images.

### Downstream Use

- Extract part-specific traits (e.g., elytra length, pronotum area) for biodiversity and phenomics research.
- Integrate segmentation outputs into the full BeetleFlow workflow alongside detection and metadata matching.

### Out-of-Scope Use

- Detection or localization of beetles in full tray images (use the BeetleFlow detection stage instead).
- Segmentation of non-beetle insects without additional fine-tuning.
- Species identification or taxonomic classification.
- Images with substantially different imaging conditions (e.g., field photos, non-museum lighting) without retraining or validation.

## Bias, Risks, and Limitations

- Models are trained on pinned ground-beetle specimens from [NEON](https://www.neonscience.org/) pitfall-trap imagery; performance may degrade on other beetle families, imaging setups, or specimen preparations.
- The 5-class and 9-class label schemes differ in granularity; choose the checkpoint that matches your annotation protocol.
- Small or occluded body parts (e.g., legs, antennas) may be harder to segment accurately.
- Segmentation quality depends on upstream cropping; poorly cropped inputs will reduce mask quality.

### Recommendations

- Use the checkpoint (`5-class` or `9-class`) that matches your target annotation schema.
- Validate model outputs on a held-out sample from your target domain before deploying at scale.
- Cite both this model and the original [sentinel-beetles](https://huggingface.co/datasets/imageomics/sentinel-beetles) dataset when publishing results.

## How to Get Started with the Model

### Load from Hugging Face Hub

```python
import torch
import cv2
import numpy as np
from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Choose "imageomics/BeetleFlow" subfolder: "5-class" or "9-class"
model_id = "imageomics/BeetleFlow"
subfolder = "5-class"  # or "9-class"

processor = Mask2FormerImageProcessor.from_pretrained(model_id, subfolder=subfolder)
model = Mask2FormerForUniversalSegmentation.from_pretrained(model_id, subfolder=subfolder)
model.to(device).eval()

# Load image (BGR → RGB)
image_bgr = cv2.imread("beetle.jpg")
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)

inputs = processor(images=image_rgb, return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)

pred_labels = processor.post_process_semantic_segmentation(
    outputs, target_sizes=[image_rgb.shape[:2]]
)[0]
pred_map = pred_labels.cpu().numpy()
```

### Batch inference (BeetleFlow pipeline)

From the [BeetleFlow repository](https://github.com/Imageomics/BeetleFlow), after cloning and installing segmentation dependencies:

```bash
python pipeline/3_segmentation/batch_inference.py \
    --model ./checkpoints/5-class \
    --input_dir /path/to/cropped/images/dir \
    --output_dir /path/to/output/dir
```

The script searches for `single_beetles` subfolders under `--input_dir` and writes predicted masks and overlay images to `--output_dir`.

### Evaluate on a test set

```bash
python pipeline/3_segmentation/test.py \
    --model ./checkpoints/5-class \
    --input /path/to/data \
    --imgsz 512 512
```

Expect `input/test_images/` and `input/test_masks/` with paired RGB mask images.

## Training Details

### Training Data

Models are trained on manually annotated beetle part segmentation data from pinned specimen images in the [Beetles as Sentinel Taxa dataset](https://huggingface.co/datasets/imageomics/sentinel-beetles). Details of the subset used for training these models:

**5-class**
- Parts: head, pronotum, elytra, legs, antennas (+ background)
- 340 labeled beetles with RGB masks
- Split: 272 training / 68 test images

**9-class**
- Parts: head, eyes, mouthparts, pronotum, elytra, tail, legs, antennas, pin (+ background)
- 330 labeled beetles with RGB masks
- Split: 264 training / 66 test images

Training uses `train_images/` and `valid_images/` (with corresponding `train_masks/` and `valid_masks/`) under the [BeetleFlow GitHub repo](https://github.com/Imageomics/BeetleFlow) data root; this repo includes mask color conventions in the [data README](https://github.com/Imageomics/BeetleFlow/tree/main/data).

### Training Procedure

#### Preprocessing

- Resize to 512×512 (width × height).
- **Training augmentations** (Albumentations): horizontal flip (p=0.5), random brightness/contrast (p=0.25), rotation (±25°).
- **Validation augmentations:** resize and normalize only.
- **Normalization:** ADE20K mean `[123.675, 116.280, 103.530]` and std `[58.395, 57.120, 57.375]` (scaled to [0, 1]).
- RGB masks are converted to integer class labels using fixed color palettes defined in [`config.py`](https://github.com/Imageomics/BeetleFlow/blob/main/pipeline/3_segmentation/config.py).

#### Training Hyperparameters

- **Optimizer:** AdamW, learning rate `1e-4`
- **Batch size:** 10
- **Epochs:** 30
- **Image size:** 512 × 512
- **LR scheduler:** MultiStepLR (optional, `--scheduler`; milestones default `[50]`, gamma `0.1`)
- **Random seed:** 42
- **Training regime:** fp32
- **Checkpoint selection:** best validation loss (`model_loss`) and best validation mIoU (`model_iou`)

#### Speeds, Sizes, Times

Training was performed on a single NVIDIA A100 GPU (40 GB) with 8 DataLoader workers. Training time was ~0.5 GPU-hours per model.

## Evaluation

Evaluation follows [`pipeline/3_segmentation/test.py`](https://github.com/Imageomics/BeetleFlow/blob/main/pipeline/3_segmentation/test.py): predictions are compared against held-out RGB masks at 512×512 resolution using the Hugging Face `evaluate` library.

### Testing Data, Factors & Metrics

#### Testing Data

Held-out test images and RGB masks (`test_images/`, `test_masks/`) from the BeetleFlow segmentation dataset (68 images for 5-class; 66 images for 9-class). Masks use the same color-to-label mapping as training.

#### Factors

- Segmentation granularity: 5-class vs. 9-class label scheme
- Body part category (per-class IoU reported separately)

#### Metrics

- **Mean IoU (mIoU):** average Intersection-over-Union across all semantic classes; primary metric
- **Mean accuracy:** average per-class pixel accuracy across all semantic classes
- **Per-category IoU:** IoU for each semantic class (background, head, pronotum, etc.)

### Results

Evaluated on the held-out test split using [`test.py`](https://github.com/Imageomics/BeetleFlow/blob/main/pipeline/3_segmentation/test.py) at 512×512 resolution.

#### 5-class

| Metric | Value |
|--------|-------|
| Mean IoU | 0.8511 |
| Mean Accuracy | 0.9102 |

| Class | IoU |
|-------|-----|
| background | 0.9501 |
| head | 0.8364 |
| pronotum | 0.9185 |
| elytra | 0.9469 |
| legs | 0.7957 |
| antennas | 0.6593 |

#### 9-class

| Metric | Value |
|--------|-------|
| Mean IoU | 0.7738 |
| Mean Accuracy | 0.8663 |

| Class | IoU |
|-------|-----|
| background | 0.9616 |
| head | 0.8309 |
| eyes | 0.6843 |
| mouthparts | 0.6008 |
| pronotum | 0.9099 |
| elytra | 0.9397 |
| tail | 0.5349 |
| legs | 0.8539 |
| antennas | 0.7008 |
| pin | 0.7213 |

#### Summary

The 5-class model achieves higher overall mIoU (0.851) than the 9-class model (0.774), reflecting the greater difficulty of fine-grained part segmentation. Both models segment large, structurally distinct regions well (elytra, pronotum, background; IoU > 0.91 for 5-class). Smaller or thinner structures are more challenging: antennas (0.659) in the 5-class model, and mouthparts (0.601), tail (0.535), and eyes (0.684) in the 9-class model. The recommended checkpoint for deployment is `model_iou` (saved when validation mIoU is highest). Use `5-class` for coarser part groupings and `9-class` when fine-grained regions (eyes, mouthparts, tail, pin) are needed.

## Model Examination

Not applicable for this release. Validation segmentation overlays are saved during training (`valid_preds/`) for qualitative inspection.

## Environmental Impact

Carbon emissions were estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://doi.org/10.48550/arXiv.1910.09700):

- **Hardware Type:** a single NVIDIA A100 GPU (40 GB)
- **Hours used:** ~0.5 GPU-hours per model
- **Cloud Provider:** Ohio Supercomputer Center (OSC)
- **Compute Region:** United States
- **Carbon Emitted:** 0.14 kg CO2 eq. (two models)

## Technical Specifications

### Model Architecture and Objective

- **Architecture:** [Mask2Former](https://huggingface.co/docs/transformers/model_doc/mask2former) for universal segmentation (`Mask2FormerForUniversalSegmentation`)
- **Backbone:** Swin Transformer Large (pretrained on ADE20K semantic segmentation)
- **Objective:** Mask-classification loss (cross-entropy + dice + mask losses) for semantic part segmentation
- **Output:** Per-pixel class map over beetle body parts
- **Input:** RGB images; processed via `Mask2FormerImageProcessor`

### Compute Infrastructure

#### Hardware

- **Training:** NVIDIA GPU with CUDA support recommended
- **Inference:** GPU recommended; CPU inference supported but slower

#### Software

- Python 3.x
- PyTorch
- [Transformers](https://huggingface.co/docs/transformers) ≥ 4.40 (trained with 4.55.2)
- OpenCV, Albumentations, `evaluate`, tqdm
- See [`requirements_segmentation.txt`](https://github.com/Imageomics/BeetleFlow/blob/main/requirements_segmentation.txt) in the BeetleFlow repository

## Citation

If you use these models, please cite BeetleFlow and the underlying dataset.

**BibTeX (software):**

```bibtex
@software{beetleflow2026,
  author = {Liu, Fangxun and Rayeed, S M and Stevens, Samuel and East, Alyson and Chiang, Cheng Hsuan and Lee, Colin and Yi, Daniel and Yang, Junke and Naik, Tejas and Wang, Ziyi and Kilrain, Connor and Buckwalter, Elijah H. and Hou, Jiacheng and Bueno, Saul Ibaven and Wang, Shuheng and Ma, Xinyue and Liu, Yifan and Tao, Zhiyuan and Zhang, Ziheng and Sokol, Eric and Belitz, Michael and Record, Sydne and Stewart, Charles V. and Chao, Wei-Lun},
  title = {BeetleFlow: An Integrative Deep Learning Pipeline for Beetle Image Processing},
  version = {1.0.0},
  year = {2026},
  url = {https://github.com/Imageomics/BeetleFlow},
  doi = {10.5281/zenodo.21251293}
}
```

**BibTeX (paper):**

```bibtex
@inproceedings{liu2025beetleflow,
  title = {BeetleFlow: An Integrative Deep Learning Pipeline for Beetle Image Processing},
  author = {Liu, Fangxun and Rayeed, S M and Stevens, Samuel and East, Alyson and Chiang, Cheng Hsuan and Lee, Colin and Yi, Daniel and Yang, Junke and Naik, Tejas and Wang, Ziyi and Kilrain, Connor and Buckwalter, Elijah H. and Hou, Jiacheng and Bueno, Saul Ibaven and Wang, Shuheng and Ma, Xinyue and Liu, Yifan and Tao, Zhiyuan and Zhang, Ziheng and Sokol, Eric and Belitz, Michael and Record, Sydne and Stewart, Charles V. and Chao, Wei-Lun},
  booktitle = {NeurIPS 2025 Workshop for Imageomics: Discovering Biological Knowledge from Images Using AI},
  year = {2025}
}
```

**Dataset:**

```bibtex
@misc{East-beetles-2025,
  author = {Alyson East and Michael Belitz and Leah Cotton and Jacqueline Dominguez and Isabelle Betancourt and S M Rayeed and Fangxun Liu and David Carlyn and Connor Kilrain and Jiaman Wu and Chandra Earl and Hilmar Lapp and Kayla I. Perry and Charles Stewart and Matthew J. Thompson and Elizabeth G. Campolongo and Wei-Lun Chao and Eric R. Sokol and Sydne Record},
  title = {Beetles as Sentinel Taxa: Predicting drought conditions from {NEON} specimen imagery (Revision 5026be7)},
  year = {2026},
  url = {https://huggingface.co/datasets/imageomics/sentinel-beetles},
  doi = {10.57967/hf/8716},
  publisher = {Hugging Face}
}
```

## Acknowledgements

This project was in part conceived at [Funcapalooza](https://github.com/Imageomics/FuncaPalooza-2025).

This work was supported by the [Imageomics Institute](https://imageomics.org), which is funded by the US National Science Foundation's Harnessing the Data Revolution (HDR) program under [Award #2118240](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2118240) (Imageomics: A New Frontier of Biological Information Powered by Knowledge-Guided Machine Learning).

This material is based in part upon work supported by the [National Ecological Observatory Network (NEON)](https://www.neonscience.org/), a program sponsored by the U.S. National Science Foundation (NSF) and operated under cooperative agreement by Battelle.

Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.

## Glossary

- **mIoU (mean Intersection over Union):** Average IoU across all semantic classes; primary evaluation metric.
- **5-class / 9-class:** Coarse vs. fine-grained beetle part label schemes (see table at top).
- **RGB mask:** Segmentation ground truth stored as a color image where each part has a fixed RGB value (see `config.py` / dataset README).

## More Information

- Full pipeline documentation: [BeetleFlow GitHub](https://github.com/Imageomics/BeetleFlow)
- Segmentation training code: [`pipeline/3_segmentation`](https://github.com/Imageomics/BeetleFlow/tree/main/pipeline/3_segmentation)
- Segmentation dataset: [`data/`](https://github.com/Imageomics/BeetleFlow/tree/main/data)

## Model Card Authors

Fangxun Liu

## Model Card Contact

Open a [Discussion](https://huggingface.co/imageomics/BeetleFlow/discussions) on this model repository, or file an issue on [GitHub](https://github.com/Imageomics/BeetleFlow/issues).