models/sam.py

import gc
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from models import torch_device
from transformers import SamModel, SamProcessor
import utils
import cv2
from scipy import ndimage

def load_sam():
    sam_model = SamModel.from_pretrained("facebook/sam-vit-base").to(torch_device)
    sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

    sam_model_dict = dict(
        sam_model = sam_model, sam_processor = sam_processor
    )

    return sam_model_dict

# Not fully backward compatible with the previous implementation
# Reference: lmdv2/notebooks/gen_masked_latents_multi_object_ref_ca_loss_modular.ipynb
def sam(sam_model_dict, image, input_points=None, input_boxes=None, target_mask_shape=None, return_numpy=True):
    """target_mask_shape: (h, w)"""
    sam_model, sam_processor = sam_model_dict['sam_model'], sam_model_dict['sam_processor']
    
    if input_boxes and isinstance(input_boxes[0], tuple):
        # Convert tuple to list
        input_boxes = [list(input_box) for input_box in input_boxes]
        
    if input_boxes and input_boxes[0] and isinstance(input_boxes[0][0], tuple):
        # Convert tuple to list
        input_boxes = [[list(input_box) for input_box in input_boxes_item] for input_boxes_item in input_boxes]
    
    with torch.no_grad():
        with torch.autocast(torch_device):
            inputs = sam_processor(image, input_points=input_points, input_boxes=input_boxes, return_tensors="pt").to(torch_device)
            outputs = sam_model(**inputs)
        masks = sam_processor.image_processor.post_process_masks(
            outputs.pred_masks.cpu().float(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
        )
        conf_scores = outputs.iou_scores.cpu().numpy()[0,0]
        del inputs, outputs
    
    gc.collect()
    torch.cuda.empty_cache()
    
    if return_numpy:
        masks = [F.interpolate(masks_item.type(torch.float), target_mask_shape, mode='bilinear').type(torch.bool).numpy() for masks_item in masks]
    else:
        masks = [F.interpolate(masks_item.type(torch.float), target_mask_shape, mode='bilinear').type(torch.bool) for masks_item in masks]

    return masks, conf_scores

def sam_point_input(sam_model_dict, image, input_points, **kwargs):
    return sam(sam_model_dict, image, input_points=input_points, **kwargs)
    
def sam_box_input(sam_model_dict, image, input_boxes, **kwargs):
    return sam(sam_model_dict, image, input_boxes=input_boxes, **kwargs)

def get_iou_with_resize(mask, masks, masks_shape):
    masks = np.array([cv2.resize(mask.astype(np.uint8) * 255, masks_shape[::-1], cv2.INTER_LINEAR).astype(bool) for mask in masks])
    return utils.iou(mask, masks)

def select_mask(masks, conf_scores, coarse_ious=None, rule="largest_over_conf", discourage_mask_below_confidence=0.85, discourage_mask_below_coarse_iou=0.2, verbose=False):
    """masks: numpy bool array"""
    mask_sizes = masks.sum(axis=(1, 2))
    
    # Another possible rule: iou with the attention mask
    if rule == "largest_over_conf":
        # Use the largest segmentation
        # Discourage selecting masks with conf too low or coarse iou is too low
        max_mask_size = np.max(mask_sizes)
        if coarse_ious is not None:
            scores = mask_sizes - (conf_scores < discourage_mask_below_confidence) * max_mask_size - (coarse_ious < discourage_mask_below_coarse_iou) * max_mask_size
        else:
            scores = mask_sizes - (conf_scores < discourage_mask_below_confidence) * max_mask_size
        if verbose:
            print(f"mask_sizes: {mask_sizes}, scores: {scores}")
    else:
        raise ValueError(f"Unknown rule: {rule}")

    mask_id = np.argmax(scores)
    mask = masks[mask_id]
    
    selection_conf = conf_scores[mask_id]
    
    if coarse_ious is not None:
        selection_coarse_iou = coarse_ious[mask_id]
    else:
        selection_coarse_iou = None

    if verbose:
        # print(f"Confidences: {conf_scores}")
        print(f"Selected a mask with confidence: {selection_conf}, coarse_iou: {selection_coarse_iou}")

    if verbose:
        plt.figure(figsize=(10, 8))
        # plt.suptitle("After SAM")
        for ind in range(3):
            plt.subplot(1, 3, ind+1)
            # This is obtained before resize.
            plt.title(f"Mask {ind}, score {scores[ind]}, conf {conf_scores[ind]:.2f}, iou {coarse_ious[ind] if coarse_ious is not None else None:.2f}")
            plt.imshow(masks[ind])
        plt.tight_layout()
        plt.show()
        plt.close()

    return mask, selection_conf

def preprocess_mask(token_attn_np_smooth, mask_th, n_erode_dilate_mask=0):
    token_attn_np_smooth_normalized = token_attn_np_smooth - token_attn_np_smooth.min()
    token_attn_np_smooth_normalized /= token_attn_np_smooth_normalized.max()
    mask_thresholded = token_attn_np_smooth_normalized > mask_th
    
    if n_erode_dilate_mask:
        mask_thresholded = ndimage.binary_erosion(mask_thresholded, iterations=n_erode_dilate_mask)
        mask_thresholded = ndimage.binary_dilation(mask_thresholded, iterations=n_erode_dilate_mask)
    
    return mask_thresholded

# The overall pipeline to refine the attention mask
def sam_refine_attn(sam_input_image, token_attn_np, model_dict, height, width, H, W, use_box_input, gaussian_sigma, mask_th_for_box, n_erode_dilate_mask_for_box, mask_th_for_point, discourage_mask_below_confidence, discourage_mask_below_coarse_iou, verbose):
    
    # token_attn_np is for visualizations
    token_attn_np_smooth = ndimage.gaussian_filter(token_attn_np, sigma=gaussian_sigma)

    # (w, h)
    mask_size_scale = height // token_attn_np_smooth.shape[1], width // token_attn_np_smooth.shape[0]

    if use_box_input:
        # box input
        mask_binary = preprocess_mask(token_attn_np_smooth, mask_th_for_box, n_erode_dilate_mask=n_erode_dilate_mask_for_box)

        input_boxes = utils.binary_mask_to_box(mask_binary, w_scale=mask_size_scale[0], h_scale=mask_size_scale[1])
        input_boxes = [input_boxes]

        masks, conf_scores = sam_box_input(model_dict, image=sam_input_image, input_boxes=input_boxes, target_mask_shape=(H, W))
    else:
        # point input
        mask_binary = preprocess_mask(token_attn_np_smooth, mask_th_for_point, n_erode_dilate_mask=0)

        # Uses the max coordinate only
        max_coord = np.unravel_index(token_attn_np_smooth.argmax(), token_attn_np_smooth.shape)
        # print("max_coord:", max_coord)
        input_points = [[[max_coord[1] * mask_size_scale[1], max_coord[0] * mask_size_scale[0]]]]

        masks, conf_scores = sam_point_input(model_dict, image=sam_input_image, input_points=input_points, target_mask_shape=(H, W))
        
    if verbose:
        plt.title("Coarse binary mask (for box for box input and for iou)")
        plt.imshow(mask_binary)
        plt.show()
    
    coarse_ious = get_iou_with_resize(mask_binary, masks, masks_shape=mask_binary.shape)

    mask_selected, conf_score_selected = select_mask(masks, conf_scores, coarse_ious=coarse_ious, 
                                                         rule="largest_over_conf", 
                                                         discourage_mask_below_confidence=discourage_mask_below_confidence, 
                                                         discourage_mask_below_coarse_iou=discourage_mask_below_coarse_iou,
                                                         verbose=True)

    return mask_selected, conf_score_selected

def sam_refine_box(sam_input_image, box, *args, **kwargs):
    # One image with one box
    sam_input_images, boxes = [sam_input_image], [[box]]
    mask_selected_batched_list, conf_score_selected_batched_list = sam_refine_boxes(sam_input_images, boxes, *args, **kwargs)
    return mask_selected_batched_list[0][0], conf_score_selected_batched_list[0][0]

def sam_refine_boxes(sam_input_images, boxes, model_dict, height, width, H, W, discourage_mask_below_confidence, discourage_mask_below_coarse_iou, verbose):
    # (w, h)
    input_boxes = [[utils.scale_proportion(box, H=height, W=width) for box in boxes_item] for boxes_item in boxes]

    masks, conf_scores = sam_box_input(model_dict, image=sam_input_images, input_boxes=input_boxes, target_mask_shape=(H, W))
    
    mask_selected_batched_list, conf_score_selected_batched_list = [], []
    
    for boxes_item, masks_item in zip(boxes, masks):
        mask_selected_list, conf_score_selected_list = [], []
        for box, three_masks in zip(boxes_item, masks_item):
            mask_binary = utils.proportion_to_mask(box, H, W, return_np=True)
            if verbose:
                # Also the box is the input for SAM
                plt.title("Binary mask from input box (for iou)")
                plt.imshow(mask_binary)
                plt.show()
                        
            coarse_ious = get_iou_with_resize(mask_binary, three_masks, masks_shape=mask_binary.shape)

            mask_selected, conf_score_selected = select_mask(three_masks, conf_scores, coarse_ious=coarse_ious, 
                                                                rule="largest_over_conf", 
                                                                discourage_mask_below_confidence=discourage_mask_below_confidence, 
                                                                discourage_mask_below_coarse_iou=discourage_mask_below_coarse_iou,
                                                                verbose=True)

            mask_selected_list.append(mask_selected)
            conf_score_selected_list.append(conf_score_selected)
        mask_selected_batched_list.append(mask_selected_list)
        conf_score_selected_batched_list.append(conf_score_selected_list)
    
    return mask_selected_batched_list, conf_score_selected_batched_list