segment-everything 0.1.0__py3-none-any.whl

segment_everything/__init__.py

@@ -0,0 +1,5 @@
+"""Utilities for rendering and training prompt-based and overlapping segmentation models (like SAM). """
+
+__version__ = '0.1.0'
+__author__ = "Brian Northan"
+__email__ = "bnorthan@gmail.com"

segment_everything/detect_and_segment.py

@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+from segment_everything.weights_helper import create_sam_model
+from segment_everything.stacked_labels import StackedLabels
+from segment_everything.vendored.mobilesamv2 import SamPredictor as SamPredictorV2
+from segment_everything.weights_helper import get_device
+from typing import Any, Generator, List
+import torch
+import os
+from segment_anything.utils.amg import calculate_stability_score
+import gc
+
+current_dir = os.path.dirname(__file__)
+
+def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
+    assert len(args) > 0 and all(
+        len(a) == len(args[0]) for a in args
+    ), "Batched iteration must have inputs of all the same size."
+    n_batches = len(args[0]) // batch_size + int(
+        len(args[0]) % batch_size != 0
+    )
+    for b in range(n_batches):
+        yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]
+
+def segment_from_stacked_labels(stacked_labels, model_type, device=None):
+    """ given stacked labels and a model re-segment all masks by calling sam on each bounding box
+
+        this function is useful when we have stacked labels with bounding boxes only (or other non-optimal masks)
+        and we want to resegment the masks using sam
+
+    Args:
+        stacked_labels (StackedLabels): input stacked labels 
+        model (SAM model): sam model 
+        device (string): device type
+
+    Returns:
+        StackedLabels: new stacked labels with resegmented labels 
+    """
+    if device is None:
+        device = get_device()
+    model = create_sam_model(model_type, device)
+    bbox_array = stacked_labels.get_bbox_np()
+    sam_masks = segment_from_bbox(stacked_labels.image, bbox_array, model, device)
+
+    return StackedLabels(sam_masks, stacked_labels.image)
+
+def segment_from_bbox(img, bounding_boxes, model, device):
+    """
+    Segments everything given the bounding boxes of the objects and the mobileSAMv2 prediction model.
+    Code from mobileSAMv2
+    """
+    predictor = SamPredictorV2(model)
+    predictor.set_image(img)
+    
+    input_boxes = predictor.transform.apply_boxes(
+        bounding_boxes, predictor.original_size
+    )  # Does this need to be transformed?
+    if device == "cuda":
+        input_boxes = torch.from_numpy(input_boxes).cuda()
+    elif device == "cpu":
+        input_boxes = torch.from_numpy(input_boxes)
+    sam_mask = []
+
+    predicted_ious = []
+    stability_scores = []
+
+    image_embedding = predictor.features
+    image_embedding = torch.repeat_interleave(image_embedding, 400, dim=0)
+
+    prompt_embedding = model.prompt_encoder.get_dense_pe()
+    prompt_embedding = torch.repeat_interleave(prompt_embedding, 400, dim=0)
+
+    for (boxes,) in batch_iterator(200, input_boxes):
+        with torch.no_grad():
+            image_embedding = image_embedding[0 : boxes.shape[0], :, :, :]
+            prompt_embedding = prompt_embedding[0 : boxes.shape[0], :, :, :]
+            sparse_embeddings, dense_embeddings = model.prompt_encoder(
+                points=None,
+                boxes=boxes,
+                masks=None,
+            )
+            low_res_masks, pred_ious = model.mask_decoder(
+                image_embeddings=image_embedding,
+                image_pe=prompt_embedding,
+                sparse_prompt_embeddings=sparse_embeddings,
+                dense_prompt_embeddings=dense_embeddings,
+                multimask_output=False,
+                simple_type=True,
+            )
+            low_res_masks = predictor.model.postprocess_masks(
+                low_res_masks, predictor.input_size, predictor.original_size
+            )
+            model.threshold_offset = 1
+            stability_score = (
+                calculate_stability_score(
+                    low_res_masks,
+                    model.mask_threshold,
+                    model.threshold_offset,
+                )
+                .cpu()
+                .numpy()
+            )
+            sam_mask_pre = (low_res_masks > model.mask_threshold) * 1.0
+            sam_mask.append(sam_mask_pre.squeeze(1))
+            predicted_ious.extend(pred_ious.cpu().numpy().flatten().tolist())
+            stability_scores.extend(stability_score.flatten().tolist())
+
+    sam_mask = torch.cat(sam_mask)
+    # predicted_ious = pred_ious.cpu().numpy()
+    cpu_segmentations = sam_mask.cpu().numpy()
+    del sam_mask
+
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    curr_anns = []
+    for idx in range(len(cpu_segmentations)):
+        ann = {
+            "segmentation": cpu_segmentations[idx],
+            "area": sum(sum(cpu_segmentations[idx])),
+            "predicted_iou": predicted_ious[idx],
+            "stability_score": stability_scores[idx],
+            "prompt_bbox": bounding_boxes[idx],
+        }
+        if (
+            cpu_segmentations[idx].max() < 1
+        ):  # this means that bboxes won't always == segmentations
+            continue
+        curr_anns.append(ann)
+    return curr_anns

segment_everything/napari_helper.py

@@ -0,0 +1,15 @@
+import napari
+from segment_everything.stacked_labels import StackedLabels
+from  napari_segment_everything import segment_everything
+
+def stacked_labels_to_napari(stacked_labels):
+    viewer = napari.Viewer()
+    segment_everything_widget=segment_everything.NapariSegmentEverything(viewer)
+    viewer.window.add_dock_widget(segment_everything_widget)
+    segment_everything_widget.load_project(stacked_labels.image, stacked_labels.mask_list)
+
+def to_napari(image, mask_list):
+    viewer = napari.Viewer()
+    segment_everything_widget=segment_everything.NapariSegmentEverything(viewer)
+    viewer.window.add_dock_widget(segment_everything_widget)
+    segment_everything_widget.load_project(image, mask_list)

segment_everything/prompt_generator.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Apr 22 23:58:14 2024
+
+@author: ian
+"""
+import cv2
+import os, sys
+from torchvision.models.detection import fasterrcnn_mobilenet_v3_large_fpn
+from torchvision.transforms import ToTensor
+from torchvision.ops import nms
+import torch
+
+from segment_everything.vendored.get_object_aware import get_object_aware_model
+
+class BaseDetector:
+    def __init__(self, model_path, trainable=False):
+        self.model_path = model_path
+        self.model_name = self.model_path.split("/")[-1]
+        self.trainable = trainable
+
+    def train(self, training_data):
+        raise NotImplementedError()
+
+    def predict(self, image_data):
+        raise NotImplementedError()
+
+
+class YoloDetector(BaseDetector):
+    def __init__(self, model_path, model_type, device, trainable=False):
+        super().__init__(model_path)
+        self.model_type = model_type
+
+        if (model_type == "ObjectAwareModelFromMobileSamV2"):
+            self.model = get_object_aware_model(model_path) #ObjectAwareModel(model_path)
+        else:
+            from ultralytics import YOLO
+            self.model = YOLO(model_path)
+
+        self.device = device
+
+    def train(self):
+        print(
+            "YOLO detector is not yet trainable, use RcnnDetector for training"
+        )
+
+    def get_results(
+        self,
+        image_data,
+        retina_masks=True,
+        imgsz=1024,
+        conf=0.4,
+        iou=0.9,
+        max_det=400,
+    ):
+        """
+        Runs YOLO and returns the YOLO results 
+
+        Parameters
+        ----------
+        image : numpy.ndarray
+            A 2D-image in grayscale or RGB.
+        imgsz : INT, optional
+            Size of the input image. The default is 1024.
+        conf : FLOAT, optional
+            Confidence threshold for the bounding boxes. Lower means more boxes will be detected. The default is 0.4.
+        iou : FLOAT, optional
+            Threshold for how many intersecting bounding boxes should be allowed. Lower means fewer intersecting boxes will be returned. The default is 0.9.
+        max_det : INT, optional
+            Maximum number of detections that will be returned. The default is 400.
+
+        Returns
+        -------
+        obj_results: YOLO results objects
+        """
+
+        image_cv2 = cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)
+
+        obj_results = self.model.predict(
+            image_cv2,
+            device=self.device,
+            retina_masks=True,
+            imgsz=imgsz,
+            conf=conf,
+            iou=iou,
+            max_det=max_det,
+        )
+
+        return obj_results
+    
+    def get_bounding_boxes(
+        self,
+        image_data,
+        retina_masks=True,
+        imgsz=1024,
+        conf=0.4,
+        iou=0.9,
+        max_det=400,
+    ):
+        """
+        Generates a series of bounding boxes in xyxy-format from an image, using the YOLOv8 ObjectAwareModel.
+
+        Parameters
+        ----------
+        image : numpy.ndarray
+            A 2D-image in grayscale or RGB.
+        imgsz : INT, optional
+            Size of the input image. The default is 1024.
+        conf : FLOAT, optional
+            Confidence threshold for the bounding boxes. Lower means more boxes will be detected. The default is 0.4.
+        iou : FLOAT, optional
+            Threshold for how many intersecting bounding boxes should be allowed. Lower means fewer intersecting boxes will be returned. The default is 0.9.
+        max_det : INT, optional
+            Maximum number of detections that will be returned. The default is 400.
+
+        Returns
+        -------
+        bounding_boxes : numpy.ndarray
+            An array of boxes in xyxy-coordinates.
+        """
+
+        print("Predicting bounding boxes for image data")
+        obj_results = self.get_results(image_data, retina_masks, imgsz, conf, iou, max_det) 
+
+        return obj_results[0].boxes.xyxy.cpu().numpy()
+  
+    def __str__(self):
+        s = f"\n{'Model':<10}: {self.model_name}\n"
+        s += f"{'Type':<10}: {str(self.model_type)}\n"
+        s += f"{'Trainable':<10}: {str(self.trainable)}"
+        return s
+
+
+class RcnnDetector(BaseDetector):
+    def __init__(self, model_path, device, trainable=True):
+        super().__init__(model_path, trainable)
+        self.model_type = "FasterRCNN"
+        if device == "mps":
+            device = "cpu"
+        self.device = device
+        self.model = fasterrcnn_mobilenet_v3_large_fpn(
+            box_detections_per_img=500,
+        ).to(self.device)
+        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
+
+    def train(self, training_data):
+        if self.trainable:
+            print("Training model")
+            print(self.model_path)
+            print(training_data)
+
+    def _get_transform(self, train):
+        from torchvision.transforms import v2 as T
+
+        transforms = []
+        if train:
+            transforms.append(T.RandomHorizontalFlip(0.5))
+        transforms.append(T.ToDtype(torch.float, scale=True))
+        transforms.append(T.ToPureTensor())
+        return T.Compose(transforms)
+
+    @torch.inference_mode()
+    def get_bounding_boxes(self, image_data, conf=0.5, iou=0.2):
+        image_cv2 = cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)
+
+        print("Predicting bounding boxes for image data")
+        convert_tensor = ToTensor()
+        eval_transform = self._get_transform(train=False)
+        tensor_image = convert_tensor(image_cv2)
+        x = eval_transform(tensor_image)
+        # convert RGBA -> RGB and move to device
+        x = x[:3, ...].to(self.device)
+        self.model.eval()
+        predictions = self.model([x])
+        pred = predictions[0]
+        #    print(pred)
+        idx_after = nms(pred["boxes"], pred["scores"], iou_threshold=iou)
+        pred_boxes = pred["boxes"][idx_after]
+        pred_scores = pred["scores"][idx_after]
+        pred_boxes_conf = pred_boxes[pred_scores > conf]
+        return pred_boxes_conf.cpu().numpy()
+
+    def __str__(self):
+        s = f"\n{'Model':<10}: {self.model_name}\n"
+        s += f"{'Type':<10}: {str(self.model_type)}\n"
+        s += f"{'Trainable':<10}: {str(self.trainable)}"
+        return s

segment_everything/py.typed

@@ -0,0 +1,5 @@
+You may remove this file if you don't intend to add types to your package
+
+Details at:
+
+https://mypy.readthedocs.io/en/stable/installed_packages.html#creating-pep-561-compatible-packages

segment_everything/stacked_label_dataset.py

@@ -0,0 +1,113 @@
+from torch.utils.data import Dataset
+import random
+import numpy as np
+
+class StackedLabelDataset(Dataset):
+    """
+    This class is used to create a dataset with stacked labels.  For an m by n image each label image is also
+    m by n array and contains only one label.   This way overlapping labels can be handled.  
+    """
+    def __init__(self, data, processor):
+        """ initializes the StackedLabelDataset
+
+        Args:
+            data (list): A list of dictionary objects.  Each dictionary object contains an image and a label collection
+            processor (_type_): _description_
+        """
+        self.data = data
+        self.processor = processor
+
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        """ returns an item from the dataset
+
+        Args:
+            idx (int): the index of the item to return
+
+        Returns:
+            dict: a dictionary containing the image and label collection
+        """
+        result = self.data[idx]
+        image = result['image']
+
+        # if number dims 2
+        if len(image.shape) == 2:
+            # add fake channels
+            image = np.stack([image, image, image], axis=-1)
+
+        segmentation = result['segmentation']
+
+        if 'prompt_bbox' in result:
+            x_min, y_min, x_max, y_max = result['prompt_bbox']
+            H, W = segmentation.shape
+            x_min = max(0, x_min - np.random.randint(0, 10))
+            x_max = min(W, x_max + np.random.randint(0, 10))
+            y_min = max(0, y_min - np.random.randint(0, 10))
+            y_max = min(H, y_max + np.random.randint(0, 10))
+            box_prompt = [x_min, y_min, x_max, y_max]
+            box_prompt = [float(x) for x in box_prompt]
+            box_prompt = [[box_prompt]]
+        else:
+            box_prompt = None
+
+        if self.random_index == False:
+            point_prompt = result['point_coords']
+        else:
+            y, x = result['indexes']
+            idx = np.random.randint(len(x))
+            point_prompt = [[x[idx], y[idx]]]
+
+        '''
+            if np.sum(segmentation) == 0:
+                print('empty')
+            else:
+                print('mask')
+
+            print('centroid is', result['point_coords'])
+            print('random prompt', point_prompt)
+            print()
+        '''
+        '''
+        if len(x) > 0:
+            idx = np.random.randint(len(x))
+            prompt = [[x[idx], y[idx]]]
+        else:    
+            prompt = result['point_coords']
+        '''        
+        '''
+        number_results = len(results)
+
+        # random id for label
+        results_id = random.randint(0, 4*number_results)
+
+        # if the label is from the original labels
+        if results_id < number_results:
+            result = results[results_id]
+            segmentation = result['segmentation']
+            y, x = np.where(segmentation)
+            idx = np.random.randint(len(x))
+            prompt = [[x[idx], y[idx]]]
+        else:
+            all_labels = item['binary']
+            not_labeled = np.logical_not(all_labels)
+            y, x = np.where(not_labeled)
+            idx = np.random.randint(len(x))
+            prompt = [[x[idx], y[idx]]]
+            segmentation = np.zeros_like(all_labels, dtype=bool)
+        #print('random prompt', prompt)
+        #print()
+        '''
+        #print('image shape is',image.shape)
+        #print('prompt is',prompt)
+
+        inputs = self.processor(image, input_points=point_prompt, input_boxes = box_prompt, return_tensors="pt")
+
+        # remove batch dimension which the processor adds by default
+        inputs = {k:v.squeeze(0) for k,v in inputs.items()}
+
+        inputs['ground_truth'] = segmentation.astype('float32')
+
+        # prompt = 
+        return inputs