sam_extension/pipeline/sam.py

import functools
from dataclasses import dataclass
import PIL
from PIL.Image import Image
import numpy as np
from typing import Union, Tuple, List, Optional, Callable
from sklearn.decomposition import PCA
import supervision as sv

import torch
from torch import nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as T

from segment_anything.utils.transforms import ResizeLongestSide
from segment_anything.predictor import preprocess, postprocess_masks
from segment_anything import build_sam, load_mobile_sam

from sam_extension.utils import add_prompts_tag, get_empty_detections, transform_coords
from sam_extension.pipeline.base import Pipeline, Output
from sam_extension.pipeline.groundingdino import GroundingDinoPipeline
from sam_extension.distillation_models.sam import load_distillation_sam, load_sam
from sam_extension.distillation_models import *

ORIGINAL_SAM_IMG_SIZE: int = 1024
PIXEL_MEAN = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
PIXEL_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
PREPROCESS = functools.partial(preprocess, ORIGINAL_SAM_IMG_SIZE, PIXEL_MEAN, PIXEL_STD)
POSTPROCESS_MASKS = functools.partial(postprocess_masks, ORIGINAL_SAM_IMG_SIZE)

@dataclass(repr=True)
class SAMEncoderOutput(Output):
    features: torch.Tensor
    interm_features: List[torch.Tensor]
    original_size: Tuple
    input_size: Tuple

@dataclass(repr=True)
class SAMEncoderProcesImgOutput(Output):
    input_image: torch.Tensor
    original_size: Tuple
    input_size: Tuple

@dataclass(repr=True)
class SAMDecoderPredictOutput(Output):
    masks_np: np.ndarray
    iou_predictions_np: np.ndarray
    low_res_masks_np: np.ndarray

@dataclass(repr=True)
class SAMDecoderPredictTorchOutput(Output):
    masks: torch.Tensor
    iou_predictions: torch.Tensor
    low_res_masks: torch.Tensor


class SAMEncoderPipeline(Pipeline):
    def __init__(self,
                 encoder: nn.Module,
                 input_img_size: Tuple,
                 multi_output: bool,
                 preprocess: Callable,
                 transform: ResizeLongestSide,
                 device: str,
                 *args,
                 **kwargs):
        super(SAMEncoderPipeline, self).__init__(*args, **kwargs)
        self.encoder = encoder
        self.input_img_size = input_img_size
        self.multi_output = multi_output
        self.preprocess = preprocess
        self.transform = transform
        self.device = device
    @classmethod
    def from_pretrained(cls, ckpt_path, device='cuda', *args, **kwargs):
        if 'sam_version' not in kwargs.keys():
            sam_version = 'sam'
        else:
            sam_version = kwargs['sam_version']
        sam = load_sam(ckpt_path, sam_version, device)
        encoder = sam.image_encoder
        encoder_type = encoder.__class__.__name__
        if encoder_type in ['TinyViT', 'FasterTinyViT', 'SAMEncoderViT', 'DINOSAMViT', 'FlashVisionTransformer']:
            multi_output = False
            if encoder_type in ['FasterTinyViT', 'SAMEncoderViT', 'DINOSAMViT', 'FlashVisionTransformer']:
                input_img_size = (encoder.img_size, encoder.img_size)
                if encoder_type == 'DINOSAMViT':
                    encoder = encoder.dino
            else:
                input_img_size = (ORIGINAL_SAM_IMG_SIZE, ORIGINAL_SAM_IMG_SIZE)
        else:
            multi_output = True
            input_img_size = (ORIGINAL_SAM_IMG_SIZE, ORIGINAL_SAM_IMG_SIZE)
        if sam.adaptor is None:
            transform = ResizeLongestSide(ORIGINAL_SAM_IMG_SIZE)
            preprocess_ = functools.partial(preprocess, ORIGINAL_SAM_IMG_SIZE, PIXEL_MEAN.to(device), PIXEL_STD.to(device))
        else:
            transform = T.Compose([
                T.Resize(input_img_size),
                T.ToTensor(),
                T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
            ])
            preprocess_ = None
        pipeline = cls(encoder=encoder,
                       input_img_size=input_img_size,
                       multi_output=multi_output,
                       preprocess=preprocess_,
                       transform=transform,
                       device=device)
        del sam, encoder
        torch.cuda.empty_cache()
        return pipeline

    def process_img(self, img: Union[Image, np.ndarray]) -> SAMEncoderProcesImgOutput:
        if self.preprocess is not None:
            if isinstance(img, Image):
                img = np.uint8(img)
            input_image = self.transform.apply_image(img)
            input_image_torch = torch.as_tensor(input_image, device=self.device)
            input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
            original_size = tuple(img.shape[:2])
            input_size = tuple(input_image_torch.shape[-2:])
            input_image = F.interpolate(self.preprocess(input_image_torch), size=self.input_img_size, mode='bilinear')
        else:
            if isinstance(img, np.ndarray):
                img = PIL.Image.fromarray(img)
            original_size = (img.size[1], img.size[0])
            if original_size[0] > original_size[1]:
                input_h = 1024
                input_w = int((1024 / original_size[0]) * original_size[1])
            else:
                input_w = 1024
                input_h = int((1024 / original_size[1]) * original_size[0])
            input_size = (input_h, input_w)
            input_image = self.transform(img)[None, ...].to(self.device)
        return SAMEncoderProcesImgOutput(input_image, original_size, input_size)
    @torch.no_grad()
    def get_visual_feature(self, x: Union[torch.Tensor, Image, np.ndarray]=None, **kwargs):
        pca_rgb = PCA(n_components=3)
        if 'sam_feature' in kwargs.keys() and 'original_size' in kwargs.keys():
            sam_feature = kwargs['sam_feature']
            original_size = kwargs['original_size']
        else:
            assert x is not None, 'please give x type Union[torch.Tensor, Image, np.ndarray] !'
            sam_encoder_output = self.forward(x, **kwargs)
            sam_feature = sam_encoder_output.features
            original_size = sam_encoder_output.original_size
            assert original_size is not None, 'please give original_size!'
        sam_feature = F.interpolate(sam_feature, size=original_size, mode='bilinear').permute(0, 2, 3, 1)
        b, h, w, c = sam_feature.shape
        sam_feature = sam_feature.view(-1, c).cpu().numpy()
        sam_feature = pca_rgb.fit_transform(sam_feature)
        sam_feature = torch.Tensor(sam_feature.reshape(h, w, 3))
        min_f, _ = sam_feature.min(-1)
        max_f, _ = sam_feature.max(-1)
        sam_feature = (sam_feature - min_f[..., None]) / (max_f[..., None] - min_f[..., None])
        sam_feature = sam_feature.cpu().numpy()
        sam_feature_image = PIL.Image.fromarray((sam_feature * 255).astype(np.uint8))
        return sam_feature_image
    def forward(self, x: Union[torch.Tensor, Image, np.ndarray], **kwargs) -> SAMEncoderOutput:
        if isinstance(x, (Image, np.ndarray)):
            process_img_output = self.process_img(x)
            x = process_img_output.input_image
            original_size = process_img_output.original_size
            input_size = process_img_output.input_size
        else:
            original_size = kwargs.pop('original_size') if 'original_size' in kwargs.keys() else None
            input_size = x.shape[-2:]
        with torch.no_grad():
            if self.multi_output:
                features, interm_features = self.encoder(x, **kwargs)
            else:
                features = self.encoder(x, **kwargs)
                if self.encoder.__class__.__name__ == 'DINO':
                    features = features.permute(0, 3, 1, 2)
                interm_features = None
        return SAMEncoderOutput(features, interm_features, original_size, input_size)

class SAMDecoderPipeline(Pipeline):
    def __init__(self,
                 prompt_encoder: nn.Module,
                 mask_decoder: nn.Module,
                 adaptor: nn.Module,
                 mask_threshold: float,
                 transform: ResizeLongestSide,
                 postprocess_masks: Callable,
                 img_size: int,
                 device: str,
                 *args,
                 **kwargs):
        super(SAMDecoderPipeline, self).__init__(*args, **kwargs)
        self.prompt_encoder = prompt_encoder
        self.mask_decoder = mask_decoder
        self.adaptor = adaptor
        self.mask_threshold = mask_threshold
        self.transform = transform
        self.postprocess_masks = postprocess_masks
        self.img_size = img_size
        self.device = device
    @classmethod
    def from_pretrained(cls, ckpt_path, device='cuda', *args, **kwargs):
        if 'sam_version' not in kwargs.keys():
            sam_version = 'sam'
        else:
            sam_version = kwargs['sam_version']
        sam = load_sam(ckpt_path, sam_version, device)
        if sam.image_encoder.__class__.__name__ == 'DINOSAMViT':
            adaptor = sam.image_encoder.adaptor
        elif sam.adaptor is not None:
            adaptor = sam.adaptor
        else:
            adaptor = None
        img_size = sam.image_encoder.img_size
        prompt_encoder = sam.prompt_encoder
        mask_decoder = sam.mask_decoder
        transform = ResizeLongestSide(ORIGINAL_SAM_IMG_SIZE)
        pipeline = cls(prompt_encoder=prompt_encoder,
                       mask_decoder=mask_decoder,
                       adaptor=adaptor,
                       mask_threshold=sam.mask_threshold,
                       transform=transform,
                       postprocess_masks=POSTPROCESS_MASKS,
                       img_size=img_size,
                       device=device)
        del sam, prompt_encoder, mask_decoder
        torch.cuda.empty_cache()
        return pipeline
    def visualize_prompt(self,
                         img: Union[Image, np.ndarray],
                         des_img: Union[Image, np.ndarray] = None,
                         point_labels: Union[List[int], np.ndarray] = None,
                         point_coords: Union[List[List[int]], np.ndarray] = None,
                         boxes: Union[List[List[int]], np.ndarray] = None,
                         pil: bool = False
                         ) -> Union[Image, np.ndarray]:
        if des_img is not None:
            if isinstance(des_img, np.ndarray):
                des_shape = tuple(des_img.shape[:2])

            else:
                des_shape = (des_img.size[1], des_img.size[0])
            src_shape = (img.size[1], img.size[0])
            point_coords, boxes = transform_coords(src_shape, des_shape, point_coords, boxes)
            return add_prompts_tag(des_img, point_labels, point_coords, boxes, pil)
        else:
            return add_prompts_tag(img, point_labels, point_coords, boxes, pil)

    def visualize_results(self,
                          img: Union[Image, np.ndarray],
                          des_img: Union[Image, np.ndarray] = None,
                          sam_encoder_output: Optional[SAMEncoderOutput] = None,
                          features: Optional[torch.Tensor] = None,
                          interm_features: Optional[List[torch.Tensor]] = None,
                          original_size: Optional[Tuple] = None,
                          input_size: Optional[Tuple] = None,
                          point_coords: Optional[np.ndarray] = None,
                          point_labels: Optional[np.ndarray] = None,
                          boxes: Optional[np.ndarray] = None,
                          texts: Optional[List] = None,
                          grounding_dino_pipeline: GroundingDinoPipeline = None,
                          box_threshold: float = 0.25,
                          text_threshold: float = 0.25,
                          nms_threshold: float = 0.8,
                          detections: Optional[sv.Detections] = None,
                          multimask_output: bool = True,
                          visualize_promts: bool = True,
                          pil: bool = False):
        if isinstance(img, Image):
            img = np.uint8(img)
        if des_img is not None:
            if isinstance(des_img, np.ndarray):
                des_shape = tuple(des_img.shape[:2])
            else:
                des_shape = (des_img.size[1], des_img.size[0])
            src_shape = img.shape[:2]
            if point_coords is not None or boxes is not None:
                des_point_coords, des_boxes = transform_coords(src_shape, des_shape, point_coords, boxes)
            else:
                des_point_coords = None
                des_boxes = None
        else:
            des_point_coords = None
            des_boxes = None
            src_shape = None
            des_shape = None
        detections = get_empty_detections() if detections is None else detections
        mask_annotator = sv.MaskAnnotator()
        result_list = []
        mask_result_list = []
        mask_list = []
        if boxes is None and point_coords is None and point_labels is None and texts is None or \
                (point_coords is not None and point_labels is not None and point_coords.shape[0] != point_labels.shape[0]):
            print('no prompt given!')
            result_list.append(img)
            return result_list
        # if boxes is not None and point_coords is not None and point_labels is not None:
        #     multimask_output = False
        def get_annotated_image(mask_annotator,
                                detections,
                                img,
                                point_labels=None,
                                point_coords=None,
                                boxes=None,
                                visualize_promts=True,
                                pil=False):
            annotated_image = mask_annotator.annotate(scene=img.copy(), detections=detections)
            if visualize_promts:
                annotated_image = add_prompts_tag(annotated_image, point_labels, point_coords, boxes=boxes, pil=pil)
            else:
                if pil:
                    annotated_image = PIL.Image.fromarray(annotated_image)
            return annotated_image
        def get_masked_image(img,
                             masks,
                             pil=True):
            masked_image_list = []
            for i in range(masks.shape[0]):
                object_rgb = img * (masks[i].reshape(img.shape[0], img.shape[1], 1))
                object_rgb = object_rgb.astype(np.uint8)
                bkgd_mask = np.where(object_rgb == 0, 1, 0)
                bkgd_mask *= 255
                bkgd_mask = bkgd_mask.astype(np.uint8)
                object_rgb += bkgd_mask
                if pil:
                    masked_image_list.append(PIL.Image.fromarray(object_rgb))
                else:
                    masked_image_list.append(object_rgb)
            return masked_image_list
        def interpolate_mask(mask_np, des_shape):
            mask_tensor = torch.tensor(mask_np, dtype=torch.float32).unsqueeze(0)
            mask_interpolate = F.interpolate(mask_tensor, size=des_shape, mode='bilinear')
            mask_interpolate = (mask_interpolate+0.5).long()
            mask_np = mask_interpolate.squeeze(0).numpy().astype(bool)
            return mask_np

        if point_coords is not None and point_labels is not None:
            
            if src_shape is not None:
                point_result = self.forward(sam_encoder_output,
                                            features,
                                            interm_features,
                                            original_size,
                                            input_size,
                                            des_point_coords,
                                            point_labels)
                masks_np = interpolate_mask(point_result.masks_np, src_shape)
            else:
                point_result = self.forward(sam_encoder_output,
                                            features,
                                            interm_features,
                                            original_size,
                                            input_size,
                                            point_coords,
                                            point_labels)
                masks_np = point_result.masks_np
            if multimask_output:
                for i in range(masks_np.shape[0]):
                    detections.mask = masks_np[i][None, ...]
                    mask_list.append(masks_np[i])
                    result_list.append(get_annotated_image(mask_annotator,
                                                           detections,
                                                           img,
                                                           point_labels=point_labels,
                                                           point_coords=point_coords,
                                                           visualize_promts=visualize_promts,
                                                           pil=pil))
                    mask_result_list += get_masked_image(img,
                                                         detections.mask,
                                                         pil=pil)
            else:
                index = np.argmax(point_result.iou_predictions_np)
                detections.mask = masks_np[index][None, ...]
                mask_list.append(masks_np[index])
                result_list.append(get_annotated_image(mask_annotator,
                                                       detections,
                                                       img,
                                                       point_labels=point_labels,
                                                       point_coords=point_coords,
                                                       visualize_promts=visualize_promts,
                                                       pil=pil))
                mask_result_list += get_masked_image(img,
                                                     detections.mask,
                                                     pil=pil)

        if boxes is not None:
            result_masks = []
            if src_shape is not None:
                boxes_ = des_boxes
            else:
                boxes_ = boxes
            if boxes_.shape[0] > 1:
                for i in range(len(boxes)):
                    box_result = self.forward(sam_encoder_output,
                                            features,
                                            interm_features,
                                            original_size,
                                            input_size,
                                            box=boxes_[i])
                    index = np.argmax(box_result.iou_predictions_np)
                    result_masks.append(box_result.masks_np[index])
                mask = np.array(result_masks)
                if src_shape is not None:
                    masks_np = interpolate_mask(mask, src_shape)
                else:
                    masks_np = mask
                mask_list.append(masks_np)
                detections.mask = masks_np
                result_list.append(get_annotated_image(mask_annotator,
                                                       detections,
                                                       img,
                                                       boxes=boxes,
                                                       visualize_promts=visualize_promts,
                                                       pil=pil))
                mask_result_list += get_masked_image(img,
                                                     detections.mask,
                                                     pil=pil)
            else:
                box_result = self.forward(sam_encoder_output,
                                          features,
                                          interm_features,
                                          original_size,
                                          input_size,
                                          box=boxes_)
                if src_shape is not None:
                    masks_np = interpolate_mask(box_result.masks_np, src_shape)
                else:
                    masks_np = box_result.masks_np

                if multimask_output:
                    for i in range(masks_np.shape[0]):
                        detections.mask = masks_np[i][None, ...]
                        mask_list.append(masks_np[i])
                        result_list.append(get_annotated_image(mask_annotator,
                                                               detections,
                                                               img,
                                                               boxes=boxes,
                                                               visualize_promts=visualize_promts,
                                                               pil=pil))
                        mask_result_list += get_masked_image(img,
                                                             detections.mask,
                                                             pil=pil)
                else:
                    index = np.argmax(box_result.iou_predictions_np)
                    detections.mask = masks_np[index][None, ...]
                    mask_list.append(masks_np[index])
                    result_list.append(get_annotated_image(mask_annotator, detections, img, boxes=boxes, pil=pil))
                    mask_result_list += get_masked_image(img,
                                                         detections.mask,
                                                         pil=pil)

        if texts is not None and grounding_dino_pipeline is not None:
            detections = grounding_dino_pipeline(img[:, :, ::-1], texts, box_threshold, text_threshold)
            box_annotator = sv.BoxAnnotator()
            nms_idx = torchvision.ops.nms(
                torch.from_numpy(detections.xyxy),
                torch.from_numpy(detections.confidence),
                nms_threshold
            ).numpy().tolist()

            detections.xyxy = detections.xyxy[nms_idx]
            detections.confidence = detections.confidence[nms_idx]
            detections.class_id = detections.class_id[nms_idx]
            labels = [
                f"{texts[class_id]} {confidence:0.2f}"
                for _, _, confidence, class_id, _
                in detections]
            result_masks = []
            if src_shape is not None:
                _, boxes_ = transform_coords(src_shape, des_shape, boxes=detections.xyxy)
            else:
                boxes_ = detections.xyxy
            for box in boxes_:
                box_result = self.forward(sam_encoder_output,
                                          features,
                                          interm_features,
                                          original_size,
                                          input_size,
                                          box=box)
                index = np.argmax(box_result.iou_predictions_np)
                result_masks.append(box_result.masks_np[index])
            mask = np.array(result_masks)
            if src_shape is not None:
                detections.mask = interpolate_mask(mask, src_shape)
            else:
                detections.mask = mask
            for i in range(detections.mask.shape[0]):
                mask_list.append(detections.mask[i, ...])
            if visualize_promts:
                annotated_image = mask_annotator.annotate(scene=img[:, :, ::-1].copy(), detections=detections)
                annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
            else:
                annotated_image = mask_annotator.annotate(scene=img[:, :, ::-1].copy(), detections=detections)

            if pil:
                result_list.append(PIL.Image.fromarray(annotated_image[:, :, ::-1]))
            else:
                result_list.append(annotated_image[:, :, ::-1])
            mask_result_list += get_masked_image(img,
                                                 detections.mask,
                                                 pil=pil)

        return result_list, mask_result_list, mask_list

    def predict(
            self,
            features: torch.Tensor,
            interm_features: List[torch.Tensor],
            original_size: Tuple,
            input_size: Tuple,
            point_coords: Optional[np.ndarray] = None,
            point_labels: Optional[np.ndarray] = None,
            box: Optional[np.ndarray] = None,
            mask_input: Optional[np.ndarray] = None,
            multimask_output: bool = True,
            return_logits: bool = False,
            hq_token_only: bool = False,
    ) -> SAMDecoderPredictOutput:
        """
        Predict masks for the given input prompts, using the currently set image.

        Arguments:
          point_coords (np.ndarray or None): A Nx2 array of point prompts to the
            model. Each point is in (X,Y) in pixels.
          point_labels (np.ndarray or None): A length N array of labels for the
            point prompts. 1 indicates a foreground point and 0 indicates a
            background point.
          box (np.ndarray or None): A length 4 array given a box prompt to the
            model, in XYXY format.
          mask_input (np.ndarray): A low resolution mask input to the model, typically
            coming from a previous prediction iteration. Has form 1xHxW, where
            for SAM, H=W=256.
          multimask_output (bool): If true, the model will return three masks.
            For ambiguous input prompts (such as a single click), this will often
            produce better masks than a single prediction. If only a single
            mask is needed, the model's predicted quality score can be used
            to select the best mask. For non-ambiguous prompts, such as multiple
            input prompts, multimask_output=False can give better results.
          return_logits (bool): If true, returns un-thresholded masks logits
            instead of a binary mask.

        Returns:
          (np.ndarray): The output masks in CxHxW format, where C is the
            number of masks, and (H, W) is the original image size.
          (np.ndarray): An array of length C containing the model's
            predictions for the quality of each mask.
          (np.ndarray): An array of shape CxHxW, where C is the number
            of masks and H=W=256. These low resolution logits can be passed to
            a subsequent iteration as mask input.
        """
        # Transform input prompts

        coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
        if point_coords is not None:
            assert (
                    point_labels is not None
            ), "point_labels must be supplied if point_coords is supplied."
            point_coords = self.transform.apply_coords(point_coords, original_size)
            coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.device)
            labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
            coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
        if box is not None:
            box = self.transform.apply_boxes(box, original_size)
            box_torch = torch.as_tensor(box, dtype=torch.float, device=self.device)
            box_torch = box_torch[None, :]
        if mask_input is not None:
            mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.device)
            mask_input_torch = mask_input_torch[None, :, :, :]

        sam_decoder_predict_torch_output = self.predict_torch(
            features,
            interm_features,
            original_size,
            input_size,
            coords_torch,
            labels_torch,
            box_torch,
            mask_input_torch,
            multimask_output,
            return_logits=return_logits,
            hq_token_only=hq_token_only,
        )

        masks_np = sam_decoder_predict_torch_output.masks[0].detach().cpu().numpy()
        iou_predictions_np = sam_decoder_predict_torch_output.iou_predictions[0].detach().cpu().numpy()
        low_res_masks_np = sam_decoder_predict_torch_output.low_res_masks[0].detach().cpu().numpy()
        return SAMDecoderPredictOutput(masks_np, iou_predictions_np, low_res_masks_np)

    @torch.no_grad()
    def predict_torch(
            self,
            features: torch.Tensor,
            interm_features: List[torch.Tensor],
            original_size: Tuple,
            input_size: Tuple,
            point_coords: Optional[torch.Tensor],
            point_labels: Optional[torch.Tensor],
            boxes: Optional[torch.Tensor] = None,
            mask_input: Optional[torch.Tensor] = None,
            multimask_output: bool = True,
            return_logits: bool = False,
            hq_token_only: bool = False,
    ) -> SAMDecoderPredictTorchOutput:
        """
        Predict masks for the given input prompts, using the currently set image.
        Input prompts are batched torch tensors and are expected to already be
        transformed to the input frame using ResizeLongestSide.

        Arguments:
          point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
            model. Each point is in (X,Y) in pixels.
          point_labels (torch.Tensor or None): A BxN array of labels for the
            point prompts. 1 indicates a foreground point and 0 indicates a
            background point.
          boxes (np.ndarray or None): A Bx4 array given a box prompt to the
            model, in XYXY format.
          mask_input (np.ndarray): A low resolution mask input to the model, typically
            coming from a previous prediction iteration. Has form Bx1xHxW, where
            for SAM, H=W=256. Masks returned by a previous iteration of the
            predict method do not need further transformation.
          multimask_output (bool): If true, the model will return three masks.
            For ambiguous input prompts (such as a single click), this will often
            produce better masks than a single prediction. If only a single
            mask is needed, the model's predicted quality score can be used
            to select the best mask. For non-ambiguous prompts, such as multiple
            input prompts, multimask_output=False can give better results.
          return_logits (bool): If true, returns un-thresholded masks logits
            instead of a binary mask.

        Returns:
          (torch.Tensor): The output masks in BxCxHxW format, where C is the
            number of masks, and (H, W) is the original image size.
          (torch.Tensor): An array of shape BxC containing the model's
            predictions for the quality of each mask.
          (torch.Tensor): An array of shape BxCxHxW, where C is the number
            of masks and H=W=256. These low res logits can be passed to
            a subsequent iteration as mask input.
        """

        if point_coords is not None:
            points = (point_coords, point_labels)
        else:
            points = None

        # Embed prompts
        sparse_embeddings, dense_embeddings = self.prompt_encoder(
            points=points,
            boxes=boxes,
            masks=mask_input,
        )

        # Predict masks
        low_res_masks, iou_predictions = self.mask_decoder(
            image_embeddings=features,
            image_pe=self.prompt_encoder.get_dense_pe(),
            sparse_prompt_embeddings=sparse_embeddings,
            dense_prompt_embeddings=dense_embeddings,
            multimask_output=multimask_output,
            hq_token_only=hq_token_only,
            interm_embeddings=interm_features,
        )

        # Upscale the masks to the original image resolution
        # masks = self.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
        masks = self.postprocess_masks(low_res_masks, input_size, original_size)

        if not return_logits:
            masks = masks > self.mask_threshold

        return SAMDecoderPredictTorchOutput(masks, iou_predictions, low_res_masks)
    def forward(self,
                sam_encoder_output: Optional[SAMEncoderOutput]=None,
                features: Optional[torch.Tensor]=None,
                interm_features: Optional[List[torch.Tensor]]=None,
                original_size: Optional[Tuple]=None,
                input_size: Optional[Tuple]=None,
                point_coords: Optional[np.ndarray] = None,
                point_labels: Optional[np.ndarray] = None,
                box: Optional[np.ndarray] = None,
                mask_input: Optional[np.ndarray] = None,
                multimask_output: bool = True,
                return_logits: bool = False,
                hq_token_only: bool = False,
                dino: bool = False
    ) -> SAMDecoderPredictOutput:
        assert sam_encoder_output or (features is not None and original_size is not None and input_size is not None), 'one of sam_encoder_output and four necessary inputs must be given!'
        if sam_encoder_output:
            features = sam_encoder_output.features
            interm_features = sam_encoder_output.interm_features
            original_size = sam_encoder_output.original_size
            input_size = sam_encoder_output.input_size
        if self.adaptor is not None:
            if dino:
                features = F.interpolate(F.normalize(features, dim=1), size=(64, 64), mode='bilinear').permute(0, 2, 3, 1)
            features = self.adaptor(features)
            #
            # else:
            #     features = self.adaptor(features, original_size)

        return self.predict(features,
                            interm_features,
                            original_size,
                            input_size,
                            point_coords,
                            point_labels,
                            box,
                            mask_input,
                            multimask_output,
                            return_logits,
                            hq_token_only)
    
'''
class SAMPipeline(Pipeline):
    @classmethod
    def from_pretrained(cls, ckpt_path, device='cuda', *args, **kwargs):
        sam_encoder_pipeline = SAMEncoderPipeline(ckpt_path, device, *args, **kwargs)
        sam_decoder_pipeline = SAMDecoderPipeline(ckpt_path, device, *args, **kwargs)
        pipeline = cls(**dict(sam_encoder_pipeline=sam_encoder_pipeline,
                              sam_decoder_pipeline=sam_decoder_pipeline,
                              device=device))
        return pipeline
'''