segment-animals 1.0.0__py3-none-any.whl

segment_animals/__init__.py

@@ -0,0 +1,53 @@
+from typing import Literal
+from PIL import Image
+import torch
+from .detect import DetectionModel, DetectionModelNames
+from .segment import SegmentationModel, SegmentationModelNames
+
+
+def get_default_device() -> Literal["cpu", "cuda", "mps"]:
+    """
+    Get the default device for model inference.
+    This function checks for CUDA availability first, then MPS (for Apple Silicon),
+    and defaults to CPU if neither is available.
+    """
+    if torch.cuda.is_available():
+        return "cuda"
+    elif "mps" in dir(torch.backends) and torch.backends.mps.is_available():
+        return "mps"
+    else:
+        return "cpu"
+
+
+
+class AutoAnimalSegmenter:
+    """
+    AutoAnimalSegmenter is a class that combines detection and segmentation models
+    to automatically detect and segment animals in images.
+    """
+
+    def __init__(
+        self,
+        detection_model_name: DetectionModelNames = "MDV5A",
+        detection_threshold: float = 0.15,
+        segmentation_model_name: SegmentationModelNames = "vit_h",
+        segmentation_device: Literal["cpu", "cuda", "mps"] = get_default_device(),
+    ):
+        self.detector = DetectionModel(
+            model_name=detection_model_name, threshold=detection_threshold
+        )
+        self.segmentor = SegmentationModel(
+            model_name=segmentation_model_name, device=segmentation_device
+        )
+
+    def process_image(self, image: Image.Image) -> tuple:
+
+        detections = self.detector.detect(image)
+        masks = self.segmentor.segment(image, detections)
+        return detections, masks
+
+
+def main() -> None:
+    print(
+        "No main function implemented yet. Use AutoAnimalSegmenter to process images."
+    )

segment_animals/detect.py

@@ -0,0 +1,57 @@
+import torch
+from megadetector.detection import run_detector
+from .models import AnimalDetection
+from typing import List, Literal
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+DetectionModelNames = Literal["MDV5A", "MDV5B"]
+
+
+class DetectionModel:
+    """
+    Model for detecting animals in images.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "MDV5A",
+        device: Literal["cpu", "cuda", "mps"] = "cpu",
+        threshold: float = 0.15,
+    ):
+        """
+        Initialize the detection model with a specified model name.
+        """
+        self.device = torch.device(device)
+        logger.info(f"Loading detection model '{model_name}' on device: {self.device}")
+
+        self.model = run_detector.load_detector(model_name)
+        logger.info(f"Model loaded: {self.model}")
+
+        self.threshold = threshold
+
+    def detect(self, image) -> List[AnimalDetection]:
+        """
+        Detect animals in the provided image.
+
+        :param image: The input image to process.
+        :return: A list of detections above a confidence threshold.
+        """
+        result = self.model.generate_detections_one_image(
+            image, image_id="", detection_threshold=self.threshold
+        )
+
+        return [
+            AnimalDetection(
+                bbox=(
+                    d["bbox"][0] * image.width,
+                    d["bbox"][1] * image.height,
+                    d["bbox"][2] * image.width,
+                    d["bbox"][3] * image.height,
+                ),
+                confidence=d["conf"],
+            )
+            for d in result["detections"]
+            if d["conf"] >= self.threshold and d["category"] == "1"
+        ]

segment_animals/model_cache.py

@@ -0,0 +1,262 @@
+"""
+Model caching and download utilities for segment_animals.
+
+Handles downloading and caching of model weights with hash verification.
+"""
+
+import os
+import hashlib
+import requests
+from pathlib import Path
+from typing import Optional, Dict, List
+
+
+# Model configurations with their download URLs, expected SHA256 hashes, and namespaces
+MODEL_CONFIGS = {
+    "vit_h": {
+        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
+        "filename": "sam_vit_h_4b8939.pth",
+        "sha256": "a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
+        "namespace": "sam",
+    },
+    "vit_l": {
+        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
+        "filename": "sam_vit_l_0b3195.pth",
+        "sha256": "3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622",
+        "namespace": "sam",
+    },
+    "vit_b": {
+        "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+        "filename": "sam_vit_b_01ec64.pth",
+        "sha256": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912",
+        "namespace": "sam",
+    },
+}
+
+
+def get_cache_dir(model_namespace: str = "sam") -> Path:
+    """
+    Get the cache directory for a specific model namespace within segment_animals.
+
+    Args:
+        model_namespace: The namespace for the models (e.g., 'sam', 'yolo', etc.)
+
+    Returns:
+        Path to the cache directory for the specified namespace
+    """
+    if os.name == "nt":  # Windows
+        cache_base = os.environ.get(
+            "LOCALAPPDATA", os.path.expanduser("~\\AppData\\Local")
+        )
+    else:  # macOS and Linux
+        cache_base = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache"))
+
+    cache_dir = Path(cache_base) / "segment_animals" / "models" / model_namespace
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def calculate_file_hash(file_path: Path) -> str:
+    """
+    Calculate SHA256 hash of a file.
+
+    Args:
+        file_path: Path to the file
+
+    Returns:
+        Hexadecimal SHA256 hash string
+    """
+    sha256_hash = hashlib.sha256()
+    with open(file_path, "rb") as f:
+        # Read file in chunks to handle large files efficiently
+        for chunk in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(chunk)
+    return sha256_hash.hexdigest()
+
+
+def verify_file_hash(file_path: Path, expected_hash: str) -> bool:
+    """
+    Verify that a file matches the expected SHA256 hash.
+
+    Args:
+        file_path: Path to the file to verify
+        expected_hash: Expected SHA256 hash
+
+    Returns:
+        True if hash matches, False otherwise
+    """
+    if not file_path.exists():
+        return False
+
+    actual_hash = calculate_file_hash(file_path)
+    return actual_hash.lower() == expected_hash.lower()
+
+
+def download_file_with_progress(url: str, destination: Path) -> None:
+    """
+    Download a file with progress indication.
+
+    Args:
+        url: URL to download from
+        destination: Path where to save the file
+    """
+    print(f"Downloading from {url}...")
+
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    total_size = int(response.headers.get("content-length", 0))
+    downloaded_size = 0
+
+    with open(destination, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            if chunk:
+                f.write(chunk)
+                downloaded_size += len(chunk)
+
+                if total_size > 0:
+                    progress = (downloaded_size / total_size) * 100
+                    print(f"\rProgress: {progress:.1f}%", end="", flush=True)
+
+    if total_size > 0:
+        print()  # New line after progress
+    print(f"Download completed: {destination}")
+
+
+def get_model_path(model_name: str, auto_download: bool = True) -> Path:
+    """
+    Get the path to a cached model, downloading it if necessary.
+
+    Args:
+        model_name: Name of the model (e.g., 'sam_vit_h_4b8939')
+        auto_download: Whether to automatically download the model if not found
+
+    Returns:
+        Path to the model file
+
+    Raises:
+        ValueError: If model_name is not recognized
+        FileNotFoundError: If model is not cached and auto_download is False
+        RuntimeError: If download fails or hash verification fails
+    """
+    if model_name not in MODEL_CONFIGS:
+        raise ValueError(
+            f"Unknown model: {model_name}. Available models: {list(MODEL_CONFIGS.keys())}"
+        )
+
+    config = MODEL_CONFIGS[model_name]
+    model_namespace = config.get(
+        "namespace", "sam"
+    )  # Default to 'sam' for backward compatibility
+    cache_dir = get_cache_dir(model_namespace)
+    model_path = cache_dir / config["filename"]
+
+    # Check if model exists and has correct hash
+    if model_path.exists():
+        print(f"Found cached model: {model_path}")
+        if verify_file_hash(model_path, config["sha256"]):
+            print("Hash verification passed")
+            return model_path
+        else:
+            print("Hash verification failed, re-downloading...")
+            model_path.unlink()  # Remove corrupted file
+
+    if not auto_download:
+        raise FileNotFoundError(
+            f"Model {model_name} not found in cache and auto_download is disabled"
+        )
+
+    # Download the model
+    try:
+        print(f"Downloading {model_name} model...")
+        download_file_with_progress(config["url"], model_path)
+
+        # Verify hash after download
+        if not verify_file_hash(model_path, config["sha256"]):
+            model_path.unlink()  # Remove corrupted download
+            raise RuntimeError(
+                f"Downloaded file hash verification failed for {model_name}"
+            )
+
+        print("Hash verification passed")
+        return model_path
+
+    except Exception as e:
+        if model_path.exists():
+            model_path.unlink()  # Clean up partial download
+        raise RuntimeError(f"Failed to download {model_name}: {e}")
+
+
+def register_model(
+    model_name: str, url: str, filename: str, sha256: str, namespace: str = "sam"
+) -> None:
+    """
+    Register a new model in the model registry.
+
+    Args:
+        model_name: Unique identifier for the model
+        url: Download URL for the model
+        filename: Local filename to save the model as
+        sha256: Expected SHA256 hash of the model file
+        namespace: Model namespace (default: 'sam')
+    """
+    MODEL_CONFIGS[model_name] = {
+        "url": url,
+        "filename": filename,
+        "sha256": sha256,
+        "namespace": namespace,
+    }
+
+
+def list_models(namespace: Optional[str] = None) -> List[str]:
+    """
+    List available models, optionally filtered by namespace.
+
+    Args:
+        namespace: Optional namespace to filter by
+
+    Returns:
+        List of model names
+    """
+    if namespace is None:
+        return list(MODEL_CONFIGS.keys())
+
+    return [
+        model_name
+        for model_name, config in MODEL_CONFIGS.items()
+        if config.get("namespace", "sam") == namespace
+    ]
+
+
+def list_namespaces() -> List[str]:
+    """
+    List all available model namespaces.
+
+    Returns:
+        List of unique namespaces
+    """
+    namespaces = set()
+    for config in MODEL_CONFIGS.values():
+        namespaces.add(config.get("namespace", "sam"))
+    return sorted(list(namespaces))
+
+
+def get_model_info(model_name: str) -> Dict[str, str]:
+    """
+    Get information about a specific model.
+
+    Args:
+        model_name: Name of the model
+
+    Returns:
+        Dictionary with model information
+
+    Raises:
+        ValueError: If model_name is not recognized
+    """
+    if model_name not in MODEL_CONFIGS:
+        raise ValueError(
+            f"Unknown model: {model_name}. Available models: {list(MODEL_CONFIGS.keys())}"
+        )
+
+    return MODEL_CONFIGS[model_name].copy()

segment_animals/models.py

@@ -0,0 +1,25 @@
+from typing import Annotated, Tuple
+from pydantic import BaseModel
+
+
+class AnimalDetection(BaseModel):
+    """
+    Model representing an animal detection.
+    """
+
+    bbox: Annotated[
+        Tuple[float, float, float, float],
+        "Bounding box coordinates (x_min, y_min, x_max, y_max)",
+    ]
+    confidence: Annotated[float, "Confidence score of the detection (0.0 to 1.0)"]
+
+
+class AnimalSegment(BaseModel):
+    """
+    Model representing an animal segmentation.
+    """
+
+    mask: Annotated[
+        list[list[float]],
+        "Segmentation mask as a binary array",
+    ]

segment_animals/segment.py

@@ -0,0 +1,71 @@
+from typing import List, Literal
+from segment_anything import SamPredictor, sam_model_registry
+from PIL import Image
+import numpy as np
+import torch
+from segment_animals.models import AnimalDetection
+from segment_animals.model_cache import get_model_path
+from segment_anything.utils.transforms import ResizeLongestSide
+
+SegmentationModelNames = Literal["vit_h", "vit_l", "vit_b"]
+
+
+class SegmentationModel:
+    """
+    Model for segmenting animals in images using SAM (Segment Anything Model).
+    """
+
+    def __init__(
+        self,
+        model_name: SegmentationModelNames = "vit_h",
+        device: Literal["cpu", "cuda", "mps"] = "cpu",
+    ):
+        # Get the model path, downloading it if necessary
+        model_path = get_model_path(model_name, auto_download=True)
+
+        self.sam = sam_model_registry[model_name](checkpoint=str(model_path))
+        self.sam.to(device)
+        self.predictor = SamPredictor(self.sam)
+        self.resize_transform = ResizeLongestSide(self.sam.image_encoder.img_size)
+
+    def segment(self, image: Image.Image, detections: List[AnimalDetection]):
+        """
+        Segment animals in the provided image.
+
+        :param image: The input image to process.
+        :param detections: The detections to guide the segmentation.
+        :return: A list of segmentation masks.
+        """
+        np_image = np.array(image.convert("RGB"))
+        self.predictor.set_image(np_image)
+
+        if not detections:
+            return torch.empty((0, *np_image.shape[:2]), dtype=torch.bool)
+
+        transformed_boxes = self.resize_transform.apply_boxes_torch(
+            torch.tensor(
+                np.atleast_2d(
+                    np.array(
+                        [
+                            [
+                                d.bbox[0],
+                                d.bbox[1],
+                                d.bbox[0] + d.bbox[2],
+                                d.bbox[1] + d.bbox[3],
+                            ]
+                            for d in detections
+                        ]
+                    )
+                )
+            ),
+            np_image.shape[:2],
+        ).to(self.sam.device)
+
+        masks, _, _ = self.predictor.predict_torch(
+            point_coords=None,
+            point_labels=None,
+            boxes=transformed_boxes,
+            multimask_output=False,
+        )
+
+        return masks.cpu()

segment_animals/util.py

@@ -0,0 +1,33 @@
+from PIL import Image
+from requests import get
+from io import BytesIO
+from pathlib import Path
+
+ImageSource = str | Path | Image.Image
+
+
+def load_image(source: ImageSource) -> Image.Image:
+    """
+    Load an image from a file path or URL.
+
+    :param source: The file path or URL of the image.
+    :return: The loaded image as a PIL Image object.
+    """
+    if isinstance(source, Image.Image):
+        # If the source is already a PIL Image, return a copy of it
+        return source.copy()
+
+    elif isinstance(source, (Path, str)) and Path(source).is_file():
+        # If the source is a file path, open the image from the file
+        return Image.open(source)
+
+    elif isinstance(source, str):
+        # If the source is a URL, fetch the image from the URL
+        response = get(source)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content))
+
+    else:
+        raise ValueError(
+            "Invalid image source. Provide a file path, URL, or PIL Image object."
+        )

segment_animals/viz.py

@@ -0,0 +1,100 @@
+from matplotlib import pyplot as plt
+from matplotlib.patches import Rectangle
+import numpy as np
+from PIL import Image
+import random
+
+
+def thumbnail(image, size=512):
+    """
+    Create a thumbnail of the image with the specified size.
+    The thumbnail will maintain the aspect ratio of the original image.
+    """
+    image = image.copy()
+    image.thumbnail((size, size))
+    return image
+
+
+def random_bright_rgb_0_1():
+    """
+    Generate a random bright RGB color with values between 0 and 1.
+    The color will be bright enough to be visible against most backgrounds.
+    """
+    return tuple(random.randint(128, 255) / 255 for _ in range(3))
+
+
+def plot_detections_and_masks(image, detections, masks):
+    """
+    Plot the detections and masks on the image.
+    """
+    plt.figure(figsize=(10, 10))
+    plt.imshow(image)
+
+    for i, (detection, mask) in enumerate(zip(detections, masks)):
+        random_color = random_bright_rgb_0_1()
+
+        plt.gca().add_patch(
+            Rectangle(
+                (detection.bbox[0], detection.bbox[1]),
+                detection.bbox[2],
+                detection.bbox[3],
+                fill=False,
+                color=random_color,
+            )
+        )
+        plt.text(
+            detection.bbox[0] + 8,
+            detection.bbox[1] - 20,
+            f"Animal {i + 1}: {detection.confidence:.2f}",
+            color="white",
+            fontsize=12,
+            bbox=dict(
+                facecolor=random_color,
+                alpha=0.5,
+                edgecolor="none",
+                boxstyle="round,pad=0.3",
+            ),
+        )
+
+        # Create an overlay for the mask in the same color
+        mask_2d = mask.squeeze()  # Remove extra dimension
+        overlay = np.zeros_like(image, dtype=np.float32)
+        overlay[mask_2d] = random_color
+        plt.imshow(overlay, alpha=0.5)
+
+    plt.axis("off")
+    plt.show()
+    return plt.gcf()  # Return the figure object for further manipulation if needed
+
+
+def mask_to_bbox(mask):
+    """Convert a mask to a bounding box."""
+
+    mask_2d = mask.squeeze().numpy().astype("uint8")
+    mask_pil = Image.fromarray(mask_2d * 255, mode="L")
+    bbox = mask_pil.getbbox()
+    return bbox
+
+
+def extract_masks(image, masks, whole_image=False):
+    """Extracts the original image filtered by each detected animal mask, with background transparent."""
+
+    extracted_images = []
+    image_np = np.array(image.convert("RGBA"))
+    for mask in masks:
+        mask_2d = mask.squeeze().numpy().astype("uint8")
+        # Create an alpha channel based on the mask
+        alpha = (mask_2d * 255).astype("uint8")
+        # Stack alpha channel to image
+        rgba = image_np.copy()
+        rgba[..., 3] = alpha
+        filtered_img = Image.fromarray(rgba, mode="RGBA")
+        extracted_images.append(filtered_img)
+
+        if not whole_image:
+            # Crop the image to the bounding box of the mask
+            bbox = mask_to_bbox(mask)
+            filtered_img = filtered_img.crop(bbox)
+            extracted_images[-1] = filtered_img
+
+    return extracted_images

segment_animals-1.0.0.dist-info/METADATA

@@ -0,0 +1,75 @@
+Metadata-Version: 2.4
+Name: segment-animals
+Version: 1.0.0
+Summary: Segment (Extract) Animals from Images - Removing Background
+Author-email: Ben Evans <ben@bluechimp.io>
+Requires-Python: >=3.10
+Requires-Dist: megadetector>=5.0.29
+Requires-Dist: pillow>=11.2.1
+Requires-Dist: requests>=2.32.4
+Requires-Dist: segment-anything
+Description-Content-Type: text/markdown
+
+# Segment Animals
+
+Segment Animals is a Python package for segmenting (extracting) animals from images using deep learning models. It provides a pipeline that combines object detection and segmentation to identify and extract animals from images, making it useful for wildlife research, conservation efforts, and any application where you wish to remove the background from images containing animals.
+
+Segment Animals builds upon the [Segment Anything](https://github.com/facebookresearch/segment-anything) and [MegaDetector](https://github.com/agentmorris/MegaDetector/blob/main/getting-started.md) models.
+
+# Installation
+
+You can install Segment Animals using pip:
+
+```bash
+pip install segment-animals
+```
+
+## Usage
+
+Here's a quick example of how to use Segment Animals, for a more detailed guide refer to the [notebook](./notebook.ipynb).
+
+### Importing the library and processing an image
+
+```python
+from segment_animals import AutoAnimalSegmenter
+from segment_animals.util import load_image
+
+model = AutoAnimalSegmenter()
+
+image = load_image("path/to/your/image.jpg")
+
+detections, masks = model.process_image(image)
+print(f"Found {len(detections)} animals.")
+```
+
+### Visualizing detections and masks
+
+```python
+from segment_animals.viz import plot_detections_and_masks
+
+plot_detections_and_masks(image, detections, masks)
+```
+
+You should then see a visualisation along the lines of this ([original image from Wikipedia](https://commons.wikimedia.org/wiki/File:Camouflaged_Predator.jpg))...
+
+![Example Segmentation](./example_viz.png)
+
+### Extracting and saving masks
+
+```python
+from segment_animals.viz import extract_masks
+
+# Setting whole_image to False will return individual masks cropped to the extent
+# of the predicted masks.
+for i, mask_extract in enumerate(extract_masks(image, masks, whole_image=False)):
+    # mask_extract is a PIL Image object so you can save it or manipulate it further
+    mask_extract.save(f"animal_mask_{i}.png")
+```
+
+Resulting in something like this:
+
+![Example Mask](./example_extract.png)
+
+# Working with Segment Animals?
+
+It'd be great to hear how you're using Segment Animals! Drop me a line at Benjamin.Evans at ioz.ac.uk or open an issue on the [GitHub repository](https://github.com/bencevans/segment-animals/issues).

segment_animals-1.0.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+segment_animals/__init__.py,sha256=R6bXAh2SQ40zkEKCkkaE27Sai7qOqVgjoz4ggRU8RDw,1686
+segment_animals/detect.py,sha256=pR6xeoYAo6MccWKtxO1N5Nzy232Bb49bZdDwNb2nlP4,1691
+segment_animals/model_cache.py,sha256=QjypUGXD2ncTI7RwrsXlUVta43HDJRXydlewzHQW_XA,7847
+segment_animals/models.py,sha256=N98DcIiYXRad_Vb0hRRRh3rQbIgwdV3zQVC_44MTa3Q,581
+segment_animals/segment.py,sha256=UqdC6SW_4UR6KqK-OBgGDKrmSdOuHFJgRTzVXmSg3CE,2384
+segment_animals/util.py,sha256=v8RC-crFNOUsNFIGkXmOvbFH390xhfApiCLUGx9ToA4,1014
+segment_animals/viz.py,sha256=fzPYARppg9o74b0iZYQ5IaNClb1Jwmxygj66gNEgKzo,3075
+segment_animals-1.0.0.dist-info/METADATA,sha256=TpAHXBxjqLyXZ6zavnLazfkT5PLEuCxj4v2QBXdHA4E,2631
+segment_animals-1.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+segment_animals-1.0.0.dist-info/entry_points.txt,sha256=EL5kd7XoXzMszTmmwZKaLmXlEmezkeTWh6kgWRaUArY,57
+segment_animals-1.0.0.dist-info/RECORD,,

segment_animals-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

segment_animals-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+segment-animals = segment_animals:main