InsectSAM: Insect Segmentation and Monitoring

Summary

• Introduction
• Parameters
• Examples
• Usage for developers
• Reference
• License

Introduction

InsectSAM is a fine-tuned version of Meta AI's segment-anything model, optimized for insect segmentation and monitoring. By combining InsectSAM with GroundingDINO (an open-set object detector), this pipeline enhances the accuracy of insect biodiversity segmentation from complex backgrounds. Designed for use with the DIOPSIS camera systems, algorithms and datasets, it allows for zero-shot object detection and precise mask generation.

With its powerful features and ease of integration, this pipeline empowers researchers and developers to create innovative ecological monitoring applications, detect and isolate insects in various environments effortlessly, and advance biodiversity research.

Parameters

Inputs

• input - (image -.png|.jpg|.jpeg): The input image containing insects to segment.
• box_threshold - (number-float, optional): Confidence threshold for bounding box detection with GroundingDINO. Default is 0.3.

Output

• output - (image -.png): The output image with detected insect segmentation masks, bounding boxes, and confidence score labels overlaid on top of the original input.

Examples

input	output

Usage for developers

Please find below the details to track the information and access the code for processing the model on our platform.

Requirements

• Python 3.10 or newer
• Git LFS for the model checkpoints
• A CUDA-capable GPU is optional

Fetch LFS objects if the repository was cloned without them, then install the dependencies:

git lfs pull
pip install -r requirements.txt

Run the demo

python demo.py

Code based on AIOZ structure

import os
from pathlib import Path
from typing import Any, Union, Literal
import tempfile

import torch
import numpy as np
from PIL import Image, ImageDraw
from transformers import pipeline, AutoProcessor, AutoModelForMaskGeneration

...
def do_ai_task(
        input: Union[str, Path],
        model_storage_directory: Union[str, Path],
        device: Literal["cpu", "cuda", "gpu"] = "cpu",
        box_threshold: float = 0.3,
        *args, **kwargs) -> Any:
    """Define AI task: load model, pre-process, post-process, overlay masks"""

    device_id = 0 if device in ["cuda", "gpu"] and torch.cuda.is_available() else -1

    sam_dir = os.path.join(model_storage_directory, "sam")
    dino_dir = os.path.join(model_storage_directory, "groundingdino")

    # 1. GroundingDINO Pipeline
    gdino_pipeline = pipeline(
        "zero-shot-object-detection",
        model=dino_dir,
        device=device_id
    )

    image = Image.open(input).convert("RGB")

    dino_results = gdino_pipeline(
        image,
        candidate_labels=["insect."],
        threshold=box_threshold,
    )

    # Filter boxes based on threshold
    filtered_results = [r for r in dino_results if r["score"] >= box_threshold]

    boxes = []
    labels = []
    scores = []

    for res in filtered_results:
        b = res["box"]
        boxes.append([b["xmin"], b["ymin"], b["xmax"], b["ymax"]])
        labels.append(res["label"])
        scores.append(res["score"])

    # 2. SAM processing
    masks = []
    if len(boxes) > 0:
        sam_processor = AutoProcessor.from_pretrained(sam_dir)
        sam_model = AutoModelForMaskGeneration.from_pretrained(sam_dir)
        if device_id == 0:
            sam_model = sam_model.to("cuda")

        # Format boxes for SAM
        input_boxes = [[boxes]]

        sam_inputs = sam_processor(image, input_boxes=input_boxes, return_tensors="pt")
        if device_id == 0:
            sam_inputs = {k: v.to("cuda") for k, v in sam_inputs.items()}

        with torch.no_grad():
            sam_outputs = sam_model(**sam_inputs)

        pred_masks = sam_processor.image_processor.post_process_masks(
            sam_outputs.pred_masks.cpu(),
            sam_inputs["original_sizes"].cpu(),
            sam_inputs["reshaped_input_sizes"].cpu()
        )[0]

        for i in range(len(boxes)):
            masks.append(pred_masks[i, 0, :, :].numpy())

    # Create mask overlay
    blended_image = show_masks_on_image(image, masks, boxes, labels, scores)

    # Save to a temporary file
    fd, fp = tempfile.mkstemp(suffix=".png")
    os.close(fd)

    blended_image.save(fp)

    return fp

Reference

This repository is based on and inspired by martintomov/InsectSAM and utilizes GroundingDINO for open-set object detection. We sincerely appreciate their generosity in sharing the models and the associated RB-IBDM resources.

License

We respect and comply with the terms of the author's apache-2.0 license cited in the Reference section.