third_parts/scripts/test_sam_part.py

import os
from pycocotools import mask as mask_util
import json
import numpy as np
import cv2
from distinctipy import distinctipy
import matplotlib.pyplot as plt
from PIL import Image
from types import MethodType
import json
import random

import torch
import torchvision
from detectron2.data import MetadataCatalog
from detectron2.structures import BitMasks, PolygonMasks
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data.detection_utils import read_image

from third_parts.APE.build_ape import build_ape_predictor


def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio

def sample_points(box, mask, min_points=3, max_points=16, dense_max_points=32):
    x0, y0, w, h = box
    aspect_ratio = w / h

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j) for n in range(min_points, max_points + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
        i * j <= max_points and i * j >= min_points)
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, w, h, 50)
    width_bin = w / target_aspect_ratio[0]
    height_bin = h / target_aspect_ratio[1]

    ret_points = []
    for wi in range(target_aspect_ratio[0]):
        xi = x0 + (wi+0.5) * width_bin
        for hi in range(target_aspect_ratio[1]):
            yi = y0 + (hi+0.5) * height_bin
            if mask[int(yi), int(xi)] > 0:
                ret_points.append((xi, yi))
    
    # if len(ret_points) < min_points:
    temp_points = []
    for wi in range(int(x0), int(x0+w)):
        for hi in range(int(y0), int(y0+h)):
            if mask[int(hi), int(wi)] > 0:
                temp_points.append((wi, hi))
    if len(temp_points)//dense_max_points < 1:
        uniform_indices = list(range(0, len(temp_points)))
    else:
        uniform_indices = list(range(0, len(temp_points), len(temp_points)//dense_max_points))
    additional_points = [temp_points[uniform_idx] for uniform_idx in uniform_indices[1:-1]]
    # ret_points = [temp_points[uniform_indices[1]], temp_points[uniform_indices[2]], temp_points[uniform_indices[3]]]
    ret_points = ret_points + additional_points
    return ret_points


def mask_iou(masks, chunk_size=50, chunk_mode=False):
    masks1 = masks.unsqueeze(1).char()  # n, 1, h, w
    masks2 = masks.unsqueeze(0).char() # 1, n, h, w

    if not chunk_mode:
        intersection = (masks1 * masks2)
        union = (masks1 + masks2 - intersection).sum(-1).sum(-1)
        intersection = intersection.sum(-1).sum(-1)
        return intersection, union
    
    def chunk_mask_iou(_chunk_size=50):

        num_chunks = masks1.shape[0] // _chunk_size
        if masks1.shape[0] % _chunk_size > 0:
            num_chunks += 1
        
        row_chunks_intersection, row_chunks_union = [], []
        for row_idx in range(num_chunks):
            col_chunks_intersection, col_chunks_union = [], []
            masks1_chunk = masks1[row_idx*_chunk_size:(row_idx+1)*_chunk_size]
            for col_idx in range(num_chunks):
                masks2_chunk = masks2[:, col_idx*_chunk_size:(col_idx+1)*_chunk_size]
                try:
                    intersection = masks1_chunk * masks2_chunk
                    temp_sum = masks1_chunk + masks2_chunk
                    union = (temp_sum - intersection).sum(-1).sum(-1)
                    intersection = intersection.sum(-1).sum(-1)
                except torch.cuda.OutOfMemoryError:
                    return False, None, None
                col_chunks_intersection.append(intersection)
                col_chunks_union.append(union)
            row_chunks_intersection.append(torch.cat(col_chunks_intersection, dim=1))
            row_chunks_union.append(torch.cat(col_chunks_union, dim=1))
        intersection = torch.cat(row_chunks_intersection, dim=0)
        union = torch.cat(row_chunks_union, dim=0)
        return True, intersection, union
    
    for c_size in [chunk_size, chunk_size//2, chunk_size//4]:
        is_ok, intersection, union = chunk_mask_iou(c_size)
        if not is_ok:
            continue
        return intersection, union

def mask_iou_v2(masks1, masks2, chunk_size=50, chunk_mode=False):
    masks1 = masks1.unsqueeze(1).char() # n, 1, h, w
    masks2 = masks2.unsqueeze(0).char()  # 1, m, h, w

    if not chunk_mode:
        intersection = (masks1 * masks2)
        union = (masks1 + masks2 - intersection).sum(-1).sum(-1)
        intersection = intersection.sum(-1).sum(-1)

        return intersection, union
    
    def chunk_mask_iou(_chunk_size=50):
        num_chunks1 = masks1.shape[0] // _chunk_size
        if masks1.shape[0] % _chunk_size > 0:
            num_chunks1 += 1
        
        num_chunks2 = masks2.shape[1] // _chunk_size
        if masks2.shape[0] % _chunk_size > 0:
            num_chunks2 += 1

        row_chunks_intersection, row_chunks_union = [], []
        for row_idx in range(num_chunks1):
            col_chunks_intersection, col_chunks_union = [], []
            masks1_chunk = masks1[row_idx*_chunk_size:(row_idx+1)*_chunk_size]
            for col_idx in range(num_chunks2):
                masks2_chunk = masks2[:, col_idx*_chunk_size:(col_idx+1)*_chunk_size]
                try:
                    intersection = masks1_chunk * masks2_chunk
                    temp_sum = masks1_chunk + masks2_chunk
                    union = (temp_sum - intersection).sum(-1).sum(-1)
                    intersection = intersection.sum(-1).sum(-1)
                except torch.cuda.OutOfMemoryError:
                    return False, None, None
                col_chunks_intersection.append(intersection)
                col_chunks_union.append(union)
            row_chunks_intersection.append(torch.cat(col_chunks_intersection, dim=1))
            row_chunks_union.append(torch.cat(col_chunks_union, dim=1))
        intersection = torch.cat(row_chunks_intersection, dim=0)
        union = torch.cat(row_chunks_union, dim=0)
        return True, intersection, union
    
    for c_size in [chunk_size, chunk_size//2, chunk_size//4]:
        is_ok, intersection, union = chunk_mask_iou(c_size)
        if not is_ok:
            continue
        return intersection, union

    return intersection, union


def mask_area(masks, chunk_size=50, chunk_mode=False):
    if not chunk_mode:
        return masks.sum(-1).sum(-1)
    
    num_chunks = masks.shape[0] // chunk_size
    if masks.shape[0] % chunk_size > 0:
        num_chunks += 1

    areas = []
    for i in range(num_chunks):
        masks_i = masks[i*chunk_size:(i+1)*chunk_size]
        areas.append(masks_i.sum(-1).sum(-1))
    return torch.cat(areas, dim=0)


def run_on_image(image_file, anno_file, save_path, sam_predictor, sam_auto_mask_generator):
    if not os.path.exists(image_file):
        return None
    file_name = os.path.basename(image_file).split('.')[0]
    with open(anno_file, 'r') as f:
        json_results = json.load(f)

    sam_image = cv2.imread(image_file)
    ori_height, ori_width = sam_image.shape[:2]
    sam_image = cv2.cvtColor(sam_image, cv2.COLOR_BGR2RGB)
    
    ori_image = Image.open(image_file)
    for ins_anno in json_results:
        root_ins_id = ins_anno['ins_id']

        object_mask = ins_anno['segmentation']
        if isinstance(object_mask["counts"], list):
            object_mask = mask_util.frPyObjects(object_mask, object_mask["size"][0], object_mask["size"][1])
        root_mask = mask_util.decode(object_mask)
        root_mask = root_mask.astype(np.uint8).squeeze()
        root_mask = torch.from_numpy(root_mask).unsqueeze(0)
        root_bbox = torchvision.ops.masks_to_boxes(root_mask)

        # crop
        root_bbox = root_bbox[0].numpy().tolist()
        box_w = root_bbox[2] - root_bbox[0]
        box_h = root_bbox[3] - root_bbox[1]
        loose_box_x0 = int(root_bbox[0] - box_w // 4)
        loose_box_y0 = int(root_bbox[1] - box_h // 4)
        loose_box_x1 = int(root_bbox[2] + box_w // 4)
        loose_box_y1 = int(root_bbox[3] + box_h // 4)
        loose_box_x0 = loose_box_x0 if loose_box_x0 > 0 else 0
        loose_box_y0 = loose_box_y0 if loose_box_y0 > 0 else 0
        loose_box_x1 = loose_box_x1 if loose_box_x1 < ori_width else ori_width
        loose_box_y1 = loose_box_y1 if loose_box_y1 < ori_height else ori_height

        loose_box_w = loose_box_x1 - loose_box_x0
        loose_box_h = loose_box_y1 - loose_box_y0
        if not (loose_box_w >= box_w and loose_box_h >= box_h):
            continue

        if loose_box_w < 256:
            padded_length_w = 256 - loose_box_w
            left_padded = padded_length_w // 2
            right_padded = padded_length_w - left_padded
            if loose_box_x0 - left_padded < 0:
                right_padded = right_padded + left_padded - loose_box_x0
                left_padded = loose_box_x0
            if loose_box_x1 + right_padded > ori_width:
                left_padded = left_padded + loose_box_x1 + right_padded - ori_width
                right_padded = ori_width - loose_box_x1
            loose_box_x0 = int(loose_box_x0 - left_padded)
            loose_box_x1 = int(loose_box_x1 + right_padded)
            loose_box_x0 = loose_box_x0 if loose_box_x0 > 0 else 0
            loose_box_x1 = loose_box_x1 if loose_box_x1 < ori_width else ori_width
        if loose_box_h < 256:
            padded_length_h = 256 - loose_box_h
            top_padded = padded_length_h // 2
            bottom_padded = padded_length_h - top_padded
            if loose_box_y0 - top_padded < 0:
                bottom_padded = bottom_padded + top_padded - loose_box_y0
                top_padded = loose_box_y0
            if loose_box_y1 + bottom_padded > ori_height:
                top_padded = top_padded + loose_box_y1 + bottom_padded - ori_height
                bottom_padded = ori_height - loose_box_y1
            loose_box_y0 = int(loose_box_y0 - top_padded)
            loose_box_y1 = int(loose_box_y1 + bottom_padded)
            loose_box_y0 = loose_box_y0 if loose_box_y0 > 0 else 0
            loose_box_y1 = loose_box_y1 if loose_box_y1 < ori_height else ori_height
        
        loose_box_w = loose_box_x1 - loose_box_x0
        loose_box_h = loose_box_y1 - loose_box_y0
        if  loose_box_w > loose_box_h:
            padded_length_h = loose_box_w - loose_box_h
            top_padded = padded_length_h // 2
            bottom_padded = padded_length_h - top_padded
            if loose_box_y0 - top_padded < 0:
                bottom_padded = bottom_padded + top_padded - loose_box_y0
                top_padded = loose_box_y0
            if loose_box_y1 + bottom_padded > ori_height:
                top_padded = top_padded + loose_box_y1 + bottom_padded - ori_height
                bottom_padded = ori_height - loose_box_y1
            loose_box_y0 = int(loose_box_y0 - top_padded)
            loose_box_y1 = int(loose_box_y1 + bottom_padded)
            loose_box_y0 = loose_box_y0 if loose_box_y0 > 0 else 0
            loose_box_y1 = loose_box_y1 if loose_box_y1 < ori_height else ori_height
        elif loose_box_h > loose_box_w:
            padded_length_w = loose_box_h - loose_box_w
            left_padded = padded_length_w // 2
            right_padded = padded_length_w - left_padded
            if loose_box_x0 - left_padded < 0:
                right_padded = right_padded + left_padded - loose_box_x0
                left_padded = loose_box_x0
            if loose_box_x1 + right_padded > ori_width:
                left_padded = left_padded + loose_box_x1 + right_padded - ori_width
                right_padded = ori_width - loose_box_x1
            loose_box_x0 = int(loose_box_x0 - left_padded)
            loose_box_x1 = int(loose_box_x1 + right_padded)
            loose_box_x0 = loose_box_x0 if loose_box_x0 > 0 else 0
            loose_box_x1 = loose_box_x1 if loose_box_x1 < ori_width else ori_width
        
        loose_box_w = loose_box_x1 - loose_box_x0
        loose_box_h = loose_box_y1 - loose_box_y0
        image_patch = ori_image[loose_box_y0:loose_box_y1, loose_box_x0:loose_box_x1, :]
        ori_image_patch_h, ori_image_patch_w = image_patch.shape[:2]
        root_mask_patch = root_mask[:, loose_box_y0:loose_box_y1, loose_box_x0:loose_box_x1]
        
        # resize the long side to 1024
        if loose_box_w > loose_box_h:
            target_w = 1024
            target_h = int(loose_box_h / loose_box_w * target_w)
        else:
            target_h = 1024
            target_w = int(loose_box_w / loose_box_h * target_h)
        image_patch = cv2.resize(image_patch, dsize=(target_w, target_h), interpolation=cv2.INTER_LINEAR)
        root_mask_patch = torch.nn.functional.interpolate(root_mask_patch[None].to(torch.float32), size=(target_h, target_w), mode="bilinear")
        root_mask_patch = (root_mask_patch[0] > 0.5).to(torch.int8)

        sam_predictor.set_image(image_patch)

        # sample points and prompt SAM
        root_bbox_patch = torchvision.ops.masks_to_boxes(root_mask_patch)
        x0, y0, x1, y1 = root_bbox_patch[0].numpy().tolist()
        ret_points = sample_points([x0, y0, x1 - x0, y1 - y0], root_mask_patch[0], min_points=3, max_points=16, dense_max_points=32)
        ret_points_list = [list(point) for point in ret_points]
        point_coords = torch.tensor(ret_points_list, device=sam_predictor.device).unsqueeze(1)
        point_labels = torch.ones(size=(point_coords.shape[0], 1), dtype=torch.int, device=sam_predictor.device)

        #TODO, sam automatically prediction
        generated_annos = sam_auto_mask_generator.generate(image_patch)
        auto_sam_masks, auto_iou_scores = [], []
        for object_anno in generated_annos:
            object_mask = object_anno["segmentation"]
            if isinstance(object_mask["counts"], list):
                object_mask = mask_util.frPyObjects(object_mask, object_mask["size"][0], object_mask["size"][1])
            mask = mask_util.decode(object_mask)
            mask = mask.astype(np.uint8).squeeze()
            auto_sam_masks.append(torch.from_numpy(mask))
            auto_iou_scores.append(object_anno['predicted_iou'])
        auto_sam_masks = torch.stack(auto_sam_masks)
        auto_iou_scores = torch.as_tensor(auto_iou_scores)

        part_masks, part_masks_score, _ = sam_predictor.predict_torch(
            point_coords=point_coords,
            point_labels=point_labels,
            boxes=None,
            multimask_output=True,
        )
        batch_size, num_masks_per_input = part_masks.shape[:2]

        print(part_masks.device)

        # first round filter, by iou score
        part_masks_area = mask_area(part_masks.flatten(0, 1), chunk_size=50, chunk_mode=True)
        part_masks_area = part_masks_area.reshape(batch_size, num_masks_per_input)
        part_masks_idx = torch.argmin(part_masks_area, dim=1)
        part_masks = torch.gather(part_masks, dim=1, index=part_masks_idx)

        print(part_masks.shape)

        part_masks_score = torch.gather(part_masks_score, dim=1, index=part_masks_idx)
        part_masks = part_masks[part_masks_score > 0.9]

        print(part_masks.shape)

        auto_sam_masks = auto_sam_masks[auto_iou_scores > 0.9]
        part_masks = torch.cat([part_masks, auto_sam_masks], dim=0)
        part_masks_score = torch.cat([part_masks_score[part_masks_score > 0.9], auto_iou_scores[auto_iou_scores > 0.9]], dim=0)

        # sort by score, from high to low
        sorted_indices = sorted(range(len(part_masks)), key=lambda k: part_masks_score[k], reverse=True)
        sorted_part_masks = torch.stack([part_masks[idx] for idx in sorted_indices], dim=0)
        
        # nms
        downsampled_part_masks = torch.nn.functional.interpolate(sorted_part_masks[None], size=(target_h//4, target_w//4), mode="bilinear")
        downsampled_part_masks = (downsampled_part_masks[0] > 0.5).to(sorted_part_masks.dtype).to("cuda")

        intersection, union = mask_iou(downsampled_part_masks, chunk_size=50, chunk_mode=True)
        mask_iou_matrix = intersection / union

        num_instances = len(mask_iou_matrix)
        keep = [True] * num_instances
        for ins_i in range(num_instances):
            if not keep[ins_i]:
                continue
            for ins_j in range(ins_i, num_instances):
                if ins_j == ins_i:
                    continue
                if mask_iou_matrix[ins_i, ins_j] > 0.8:
                    keep[ins_j] = False
        
        # roc
        downsampled_root_mask_patch = torch.nn.functional.interpolate(root_mask_patch[None].to(torch.float32), size=(target_h//4, target_w//4), mode="bilinear")
        downsampled_root_mask_patch = (downsampled_root_mask_patch[0] > 0.5).to(root_mask_patch.dtype).to("cuda")

        intersection, union = mask_iou_v2(downsampled_root_mask_patch, downsampled_part_masks, chunk_size=50, chunk_mode="bilinear")
        downsampled_part_masks_area = mask_area(downsampled_part_masks, chunk_mode=True, chunk_size=50)
        mask_iou = intersection[0] / union[0]
        mask_roc = intersection[0] / downsampled_part_masks_area

        maybe_is_part = (mask_iou < 0.8) & (mask_roc > 0.95) & torch.as_tensor(keep)
        
        if not torch.any(maybe_is_part):
            continue

        left_part_masks = sorted_part_masks[maybe_is_part]
        left_part_masks = torch.nn.functional.interpolate(left_part_masks[None].to(torch.float32), size=(ori_image_patch_h, ori_image_patch_w), mode="bilinear")
        left_part_masks = (left_part_masks[0] > 0.5).to(root_mask.dtype).to(root_mask.device)
        full_size_part_masks = torch.zeros_like(root_mask).repeat(left_part_masks.shape[0], 1, 1)
        full_size_part_masks[:, loose_box_y0:loose_box_y1, loose_box_x0:loose_box_x1] = left_part_masks
        full_size_part_masks = full_size_part_masks.cpu().numpy()

        save_json_results = []
        for part_idx, mask in enumerate(full_size_part_masks):
            rle = mask_util.encode(np.array(mask[:, :, None], order="F", dtype="uint8"))[0]
            rle["counts"] = rle["counts"].decode("utf-8")
            save_json_results.append({
                "root_id": root_ins_id,
                "part_id": part_idx+1,
                "segmentation": rle,
            })

        with open(os.path.join(save_path, file_name+'.json'), 'w') as f:
            json.dump(save_json_results, f)