Upload 19 files

.gitattributes

@@ -33,3 +33,13 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/dinosaur1.png filter=lfs diff=lfs merge=lfs -text
+images/dinosaur2.png filter=lfs diff=lfs merge=lfs -text
+images/dinosaur3.png filter=lfs diff=lfs merge=lfs -text
+images/elephant1.png filter=lfs diff=lfs merge=lfs -text
+images/elephant2.png filter=lfs diff=lfs merge=lfs -text
+images/elephant3.png filter=lfs diff=lfs merge=lfs -text
+images/hmbb3.png filter=lfs diff=lfs merge=lfs -text
+images/horse1.png filter=lfs diff=lfs merge=lfs -text
+images/horse2.png filter=lfs diff=lfs merge=lfs -text
+images/horse3.png filter=lfs diff=lfs merge=lfs -text

Matcher.py

@@ -0,0 +1,452 @@
+import os
+from os import path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+
+import numpy as np
+import cv2
+import ot
+import math
+from scipy.optimize import linear_sum_assignment
+
+from segment_anything import sam_model_registry
+from segment_anything import SamAutomaticMaskGenerator
+from dinov2.models import vision_transformer as vits
+import dinov2.utils.utils as dinov2_utils
+from dinov2.data.transforms import MaybeToTensor, make_normalize_transform
+
+from matcher.k_means import kmeans_pp
+
+import random
+
+class Matcher:
+    def __init__(
+            self,
+            encoder,
+            generator=None,
+            input_size=518,
+            num_centers=8,
+            use_box=False,
+            use_points_or_centers=True,
+            sample_range=(4, 6),
+            max_sample_iterations=30,
+            alpha=1.,
+            beta=0.,
+            exp=0.,
+            score_filter_cfg=None,
+            num_merging_mask=10,
+            device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),
+    ):
+        # models
+        self.encoder = encoder
+        self.generator = generator
+        self.rps = None
+
+        if not isinstance(input_size, tuple):
+            input_size = (input_size, input_size)
+        self.input_size = input_size
+
+        # transforms for image encoder
+        self.encoder_transform = transforms.Compose([
+            MaybeToTensor(),
+            make_normalize_transform(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ])
+
+        self.tar_img = None
+        self.tar_img_np = None
+        self.tar_img_ori_size = None
+
+        self.ref_imgs = None
+        self.ref_masks_pool = None
+        self.nshot = None
+
+        self.encoder_img_size = None
+        self.encoder_feat_size = None
+
+        self.num_centers = num_centers
+        self.use_box = use_box
+        self.use_points_or_centers = use_points_or_centers
+        self.sample_range = sample_range
+        self.max_sample_iterations =max_sample_iterations
+
+        self.alpha, self.beta, self.exp = alpha, beta, exp
+        assert score_filter_cfg is not None
+        self.score_filter_cfg = score_filter_cfg
+        self.num_merging_mask = num_merging_mask
+
+        self.device = device
+
+    def set_reference(self, imgs, masks):
+
+        def reference_masks_verification(masks):
+            if masks.sum() == 0:
+                _, _, sh, sw = masks.shape
+                masks[..., (sh // 2 - 7):(sh // 2 + 7), (sw // 2 - 7):(sw // 2 + 7)] = 1
+            return masks
+
+        imgs = imgs.flatten(0, 1)  # bs, 3, h, w
+        img_size = imgs.shape[-1]
+        assert img_size == self.input_size[-1]
+        feat_size = img_size // self.encoder.patch_size
+
+        self.encoder_img_size = img_size
+        self.encoder_feat_size = feat_size
+
+        # process reference masks
+        masks = reference_masks_verification(masks)
+        masks = masks.permute(1, 0, 2, 3)  # ns, 1, h, w
+        ref_masks_pool = F.avg_pool2d(masks.float(), (self.encoder.patch_size, self.encoder.patch_size))
+        nshot = ref_masks_pool.shape[0]
+        ref_masks_pool = (ref_masks_pool > self.generator.predictor.model.mask_threshold).float()
+        ref_masks_pool = ref_masks_pool.reshape(-1)  # nshot, N
+
+        self.ref_imgs = imgs
+        self.ref_masks_pool = ref_masks_pool
+        self.nshot = nshot
+
+    def set_target(self, img, tar_img_ori_size):
+
+        img_h, img_w = img.shape[-2:]
+        assert img_h == self.input_size[0] and img_w == self.input_size[1]
+
+        # transform query to numpy as input of sam
+        img_np = img.mul(255).byte()
+        img_np = img_np.squeeze(0).permute(1, 2, 0).cpu().numpy()
+
+        self.tar_img = img
+        self.tar_img_np = img_np
+        self.tar_img_ori_size = tar_img_ori_size
+
+    def set_rps(self):
+        if self.rps is None:
+            assert self.encoder_feat_size is not None
+            self.rps = RobustPromptSampler(
+                encoder_feat_size=self.encoder_feat_size,
+                sample_range=self.sample_range,
+                max_iterations=self.max_sample_iterations
+            )
+
+
+    def predict(self):
+
+        ref_feats, tar_feat = self.extract_img_feats()
+        all_points, box, S, C, reduced_points_num = self.patch_level_matching(ref_feats=ref_feats, tar_feat=tar_feat)
+        points = self.clustering(all_points) if not self.use_points_or_centers else all_points
+        self.set_rps()
+        mask, mask_list = self.mask_generation(self.tar_img_np, points, box, all_points, self.ref_masks_pool, C)
+        return mask, mask_list
+
+
+    def extract_img_feats(self):
+
+        ref_imgs = torch.cat([self.encoder_transform(rimg)[None, ...] for rimg in self.ref_imgs], dim=0)
+        tar_img = torch.cat([self.encoder_transform(timg)[None, ...] for timg in self.tar_img], dim=0)
+
+        ref_feats = self.encoder.forward_features(ref_imgs.to(self.device))["x_prenorm"][:, 1:]
+        tar_feat = self.encoder.forward_features(tar_img.to(self.device))["x_prenorm"][:, 1:]
+        # ns, N, c = ref_feats.shape
+        ref_feats = ref_feats.reshape(-1, self.encoder.embed_dim)  # ns*N, c
+        tar_feat = tar_feat.reshape(-1, self.encoder.embed_dim)  # N, c
+
+        ref_feats = F.normalize(ref_feats, dim=1, p=2) # normalize for cosine similarity
+        tar_feat = F.normalize(tar_feat, dim=1, p=2)
+
+        return ref_feats, tar_feat
+
+    def patch_level_matching(self, ref_feats, tar_feat):
+
+        # forward matching
+        S = ref_feats @ tar_feat.t()  # ns*N, N
+        C = (1 - S) / 2  # distance
+
+        S_forward = S[self.ref_masks_pool.flatten().bool()]
+
+        indices_forward = linear_sum_assignment(S_forward.cpu(), maximize=True)
+        indices_forward = [torch.as_tensor(index, dtype=torch.int64, device=self.device) for index in indices_forward]
+        sim_scores_f = S_forward[indices_forward[0], indices_forward[1]]
+        indices_mask = self.ref_masks_pool.flatten().nonzero()[:, 0]
+
+        # reverse matching
+        S_reverse = S.t()[indices_forward[1]]
+        indices_reverse = linear_sum_assignment(S_reverse.cpu(), maximize=True)
+        indices_reverse = [torch.as_tensor(index, dtype=torch.int64, device=self.device) for index in indices_reverse]
+        retain_ind = torch.isin(indices_reverse[1], indices_mask)
+        if not (retain_ind == False).all().item():
+            indices_forward = [indices_forward[0][retain_ind], indices_forward[1][retain_ind]]
+            sim_scores_f = sim_scores_f[retain_ind]
+        inds_matched, sim_matched = indices_forward, sim_scores_f
+
+        reduced_points_num = len(sim_matched) // 2 if len(sim_matched) > 40 else len(sim_matched)
+        sim_sorted, sim_idx_sorted = torch.sort(sim_matched, descending=True)
+        sim_filter = sim_idx_sorted[:reduced_points_num]
+        points_matched_inds = indices_forward[1][sim_filter]
+
+        points_matched_inds_set = torch.tensor(list(set(points_matched_inds.cpu().tolist())))
+        points_matched_inds_set_w = points_matched_inds_set % (self.encoder_feat_size)
+        points_matched_inds_set_h = points_matched_inds_set // (self.encoder_feat_size)
+        idxs_mask_set_x = (points_matched_inds_set_w * self.encoder.patch_size + self.encoder.patch_size // 2).tolist()
+        idxs_mask_set_y = (points_matched_inds_set_h * self.encoder.patch_size + self.encoder.patch_size // 2).tolist()
+
+        ponits_matched = []
+        for x, y in zip(idxs_mask_set_x, idxs_mask_set_y):
+            if int(x) < self.input_size[1] and int(y) < self.input_size[0]:
+                ponits_matched.append([int(x), int(y)])
+        ponits = np.array(ponits_matched)
+
+        if self.use_box:
+            box = np.array([
+                max(ponits[:, 0].min(), 0),
+                max(ponits[:, 1].min(), 0),
+                min(ponits[:, 0].max(), self.input_size[1] - 1),
+                min(ponits[:, 1].max(), self.input_size[0] - 1),
+            ])
+        else:
+            box = None
+
+        return ponits, box, S, C, reduced_points_num
+
+    def clustering(self, points):
+
+        num_centers = min(self.num_centers, len(points))
+        flag = True
+        while (flag):
+            centers, cluster_assignment = kmeans_pp(points, num_centers)
+            id, fre = torch.unique(cluster_assignment, return_counts=True)
+            if id.shape[0] == num_centers:
+                flag = False
+            else:
+                print('Kmeans++ failed, re-run')
+        centers = np.array(centers).astype(np.int64)
+        return centers
+
+
+    def mask_generation(self, tar_img_np, points, box, all_ponits, ref_masks_pool, C):
+        samples_list, label_list = self.rps.sample_points(points)
+        tar_masks_ori = self.generator.generate(
+            tar_img_np,
+            select_point_coords=samples_list,
+            select_point_labels=label_list,
+            select_box=[box] if self.use_box else None,
+        )
+        tar_masks = torch.cat(
+            [torch.from_numpy(qmask['segmentation']).float()[None, None, ...].to(self.device) for
+             qmask in tar_masks_ori], dim=0).cpu().numpy() > 0
+
+        # append to original results
+        purity = torch.zeros(tar_masks.shape[0])
+        coverage = torch.zeros(tar_masks.shape[0])
+        emd = torch.zeros(tar_masks.shape[0])
+
+        samples = samples_list[-1]
+        labels = torch.ones(tar_masks.shape[0], samples.shape[1])
+        samples = torch.ones(tar_masks.shape[0], samples.shape[1], 2)
+
+        # compute scores for each mask
+        for i in range(len(tar_masks)):
+            purity_, coverage_, emd_, sample_, label_, mask_ = \
+                self.rps.get_mask_scores(
+                    points=points,
+                    masks=tar_masks[i],
+                    all_points=all_ponits,
+                    emd_cost=C,
+                    ref_masks_pool=ref_masks_pool
+                )
+            assert np.all(mask_ == tar_masks[i])
+            purity[i] = purity_
+            coverage[i] = coverage_
+            emd[i] = emd_
+
+        pred_masks = tar_masks.squeeze(1)
+        metric_preds = {
+            "purity": purity,
+            "coverage": coverage,
+            "emd": emd
+        }
+
+        scores = self.alpha * emd + self.beta * purity * coverage ** self.exp
+
+        def check_pred_mask(pred_masks):
+            if len(pred_masks.shape) < 3:  # avoid only one mask
+                pred_masks = pred_masks[None, ...]
+            return pred_masks
+
+        pred_masks = check_pred_mask(pred_masks)
+
+        # filter the false-positive mask fragments by using the proposed metrics
+        for metric in ["coverage", "emd", "purity"]:
+            if self.score_filter_cfg[metric] > 0:
+                thres = min(self.score_filter_cfg[metric], metric_preds[metric].max())
+                idx = torch.where(metric_preds[metric] >= thres)[0]
+                scores = scores[idx]
+                samples = samples[idx]
+                labels = labels[idx]
+                pred_masks = check_pred_mask(pred_masks[idx])
+
+                for key in metric_preds.keys():
+                    metric_preds[key] = metric_preds[key][idx]
+
+        #  score-based masks selection, masks merging
+        if self.score_filter_cfg["score_filter"]:
+
+            distances = 1 - scores
+            distances, rank = torch.sort(distances, descending=False)
+            distances_norm = distances - distances.min()
+            distances_norm = distances_norm / (distances.max() + 1e-6)
+            filer_dis = distances < self.score_filter_cfg["score"]
+            filer_dis[..., 0] = True
+            filer_dis_norm = distances_norm < self.score_filter_cfg["score_norm"]
+            filer_dis = filer_dis * filer_dis_norm
+
+            pred_masks = check_pred_mask(pred_masks)
+            masks = pred_masks[rank[filer_dis][:self.num_merging_mask]]
+            masks = check_pred_mask(masks)
+            mask_list = masks
+            masks = masks.sum(0) > 0
+            masks = masks[None, ...]
+
+        else:
+
+            topk = min(self.num_merging_mask, scores.size(0))
+            topk_idx = scores.topk(topk)[1]
+            topk_samples = samples[topk_idx].cpu().numpy()
+            topk_scores = scores[topk_idx].cpu().numpy()
+            topk_pred_masks = pred_masks[topk_idx]
+            topk_pred_masks = check_pred_mask(topk_pred_masks)
+
+            if self.score_filter_cfg["topk_scores_threshold"] > 0:
+                # map scores to 0-1
+                topk_scores = topk_scores / (topk_scores.max())
+
+            idx = topk_scores > self.score_filter_cfg["topk_scores_threshold"]
+            topk_samples = topk_samples[idx]
+
+            topk_pred_masks = check_pred_mask(topk_pred_masks)
+            topk_pred_masks = topk_pred_masks[idx]
+            mask_list = []
+            for i in range(len(topk_samples)):
+                mask = topk_pred_masks[i][None, ...]
+                mask_list.append(mask)
+            mask_list = np.concatenate(mask_list, axis=0)
+            masks = np.sum(mask_list, axis=0) > 0
+            masks = check_pred_mask(masks)
+
+        tar_img_ori_size = self.tar_img_ori_size
+        mask = torch.tensor(masks, device=self.device)[None, ...]
+        mask = F.interpolate(mask.float(), tar_img_ori_size, mode="bilinear", align_corners=False) > 0
+        mask = mask.squeeze(0).cpu().numpy()
+        if mask_list is not None:
+            mask_list = torch.tensor(mask_list, device=self.device)[:, None, ...]
+            mask_list = F.interpolate(mask_list.float(), tar_img_ori_size, mode="bilinear", align_corners=False)
+            mask_list = mask_list.squeeze(0).cpu().numpy()
+
+        return mask, mask_list
+
+
+    def clear(self):
+
+        self.tar_img = None
+        self.tar_img_np = None
+        self.tar_img_ori_size = None
+
+        self.ref_imgs = None
+        self.ref_masks_pool = None
+        self.nshot = None
+
+        self.encoder_img_size = None
+        self.encoder_feat_size = None
+
+
+
+class RobustPromptSampler:
+
+    def __init__(
+        self,
+        encoder_feat_size,
+        sample_range,
+        max_iterations
+    ):
+        self.encoder_feat_size = encoder_feat_size
+        self.sample_range = sample_range
+        self.max_iterations = max_iterations
+
+
+    def get_mask_scores(self, points, masks, all_points, emd_cost, ref_masks_pool):
+
+        def is_in_mask(point, mask):
+            # input: point: n*2, mask: h*w
+            # output: n*1
+            h, w = mask.shape
+            point = point.astype(np.int)
+            point = point[:, ::-1]  # y,x
+            point = np.clip(point, 0, [h - 1, w - 1])
+            return mask[point[:, 0], point[:, 1]]
+
+        ori_masks = masks
+        masks = cv2.resize(
+            masks[0].astype(np.float32),
+            (self.encoder_feat_size, self.encoder_feat_size),
+            interpolation=cv2.INTER_AREA)
+        if masks.max() <= 0:
+            thres = masks.max() - 1e-6
+        else:
+            thres = 0
+        masks = masks > thres
+
+        # 1. emd
+        emd_cost_pool = emd_cost[ref_masks_pool.flatten().bool(), :][:, masks.flatten()]
+        emd = ot.emd2(a=[1. / emd_cost_pool.shape[0] for i in range(emd_cost_pool.shape[0])],
+                      b=[1. / emd_cost_pool.shape[1] for i in range(emd_cost_pool.shape[1])],
+                      M=emd_cost_pool.cpu().numpy())
+        emd_score = 1 - emd
+
+        labels = np.ones((points.shape[0],))
+
+        # 2. purity and coverage
+        assert all_points is not None
+        points_in_mask = is_in_mask(all_points, ori_masks[0])
+        points_in_mask = all_points[points_in_mask]
+        # here we define two metrics for local matching , purity and coverage
+        # purity: points_in/mask_area, the higher means the denser points in mask
+        # coverage: points_in / all_points, the higher means the mask is more complete
+        mask_area = max(float(masks.sum()), 1.0)
+        purity = points_in_mask.shape[0] / mask_area
+        coverage = points_in_mask.shape[0] / all_points.shape[0]
+        purity = torch.tensor([purity]) + 1e-6
+        coverage = torch.tensor([coverage]) + 1e-6
+        return purity, coverage, emd_score, points, labels, ori_masks
+
+    def combinations(self, n, k):
+        if k > n:
+            return []
+        if k == 0:
+            return [[]]
+        if k == n:
+            return [[i for i in range(n)]]
+        res = []
+        for i in range(n):
+            for j in self.combinations(i, k - 1):
+                res.append(j + [i])
+        return res
+
+    def sample_points(self, points):
+        # return list of arrary
+
+        sample_list = []
+        label_list = []
+        for i in range(min(self.sample_range[0], len(points)), min(self.sample_range[1], len(points)) + 1):
+            if len(points) > 8:
+                index = [random.sample(range(len(points)), i) for j in range(self.max_iterations)]
+                sample = np.take(points, index, axis=0)  # (max_iterations * i) * 2
+            else:
+                index = self.combinations(len(points), i)
+                sample = np.take(points, index, axis=0)  # i * n * 2
+
+            # generate label  max_iterations * i
+            label = np.ones((sample.shape[0], i))
+            sample_list.append(sample)
+            label_list.append(label)
+        return sample_list, label_list
+

app_running.py

@@ -0,0 +1,173 @@
+import os
+import numpy as np
+import random
+import time
+import gradio as gr
+from gradio_demo.runner import Runner
+import matplotlib.pyplot as plt
+
+def show_mask(mask, ax, color='blue'):
+    if color == 'blue':
+        # reference, blue
+        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
+    else:
+        # target, green
+        color = np.array([78 / 255, 238 / 255, 148 / 255, 0.6])
+
+    # if random_color:
+    #     color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    # else:
+    #     color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    ax.imshow(mask_image)
+
+
+def show_points(coords, labels, ax, marker_size=375):
+    pos_points = coords[labels == 1]
+    neg_points = coords[labels == 0]
+    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
+               linewidth=1.25)
+    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
+               linewidth=1.25)
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
+
+
+def show_img_point_box_mask(img, input_point=None, input_label=None, box=None, masks=None, save_path=None, mode='mask', color='blue'):
+
+    if mode == 'point':
+        # point
+        plt.figure(figsize=(10, 10))
+        plt.imshow(img)
+        show_points(input_point, input_label, plt.gca())
+        plt.axis('on')
+        plt.savefig(save_path, bbox_inches='tight')
+    elif mode == 'box':
+        # box
+        plt.figure(figsize=(10, 10))
+        plt.imshow(img)
+        show_box(box, plt.gca())
+        plt.axis('on')
+        plt.savefig(save_path, bbox_inches='tight')
+    else:
+        # mask
+        plt.figure(figsize=(10, 10))
+        plt.imshow(img)
+        show_mask(masks, plt.gca(), color=color)
+        plt.axis('off')
+        plt.savefig(save_path, bbox_inches='tight')
+    plt.close()
+
+
+def create_oss_demo(
+        runner: Runner,
+        pipe: None = None
+) -> gr.Blocks:
+
+    examples = [
+        ['./gradio_demo/images/horse1.png', './gradio_demo/images/horse2.png', './gradio_demo/images/horse3.png'],
+        ['./gradio_demo/images/hmbb1.png', './gradio_demo/images/hmbb2.png', './gradio_demo/images/hmbb3.png'],
+        ['./gradio_demo/images/earth1.png', './gradio_demo/images/earth2.png', './gradio_demo/images/earth3.png'],
+        ['./gradio_demo/images/elephant1.png', './gradio_demo/images/elephant2.png', './gradio_demo/images/elephant3.png'],
+        ['./gradio_demo/images/dinosaur1.png', './gradio_demo/images/dinosaur2.png', './gradio_demo/images/dinosaur3.png'],
+    ]
+
+    with gr.Blocks() as oss_demo:
+        with gr.Column():
+
+            # inputs
+            with gr.Row():
+                img_input_prompt = gr.ImageMask(label='Prompt (提示图)')
+                img_input_target1 = gr.Image(label='Target 1 (测试图1)')
+                img_input_target2 = gr.Image(label='Target 2 (测试图2)')
+
+            version = gr.inputs.Radio(['version 1 (🔺 multiple instances  🔻 whole, 🔻 part)',
+                                       'version 2 (🔻 multiple instances  🔺 whole, 🔻 part)',
+                                       'version 3 (🔻 multiple instances  🔻 whole, 🔺 part)'],
+                                      type="value", default='version 1 (🔺 whole, 🔻 part)',
+                                      label='Multiple Instances (version 1), Single Instance (version 2), Part of a object (version 3)')
+
+            with gr.Row():
+                submit1 = gr.Button("提交 (Submit)")
+                clear = gr.Button("清除 (Clear)")
+            info = gr.Text(label="Processing result: ", interactive=False)
+
+            # decision
+            K = gr.Slider(0, 10, 10, step=1, label="Controllable mask output", interactive=True)
+            submit2 = gr.Button("提交 (Submit)")
+
+            # outputs
+            with gr.Row():
+                img_output_pmt = gr.Image(label='Prompt (提示图)')
+                img_output_tar1 = gr.Image(label='Output 1 (输出图1)')
+                img_output_tar2 = gr.Image(label='Output 2 (输出图2)')
+
+        # images
+        gr.Examples(
+            examples=examples,
+            fn=runner.inference_oss_ops,
+            inputs=[img_input_prompt, img_input_target1, img_input_target2],
+            outputs=info
+        )
+
+        submit1.click(
+            fn=runner.inference_oss_ops,
+            inputs=[img_input_prompt, img_input_target1, img_input_target2, version],
+            outputs=info
+        )
+        submit2.click(
+            fn=runner.controllable_mask_output,
+            inputs=K,
+            outputs=[img_output_pmt, img_output_tar1, img_output_tar2]
+        )
+
+        clear.click(
+            fn=runner.clear_fn,
+            inputs=None,
+            outputs=[img_input_prompt, img_input_target1, img_input_target2, info, img_output_pmt, img_output_tar1, img_output_tar2],
+            queue=False
+        )
+
+    return oss_demo
+
+
+def create_vos_demo(
+        runner: Runner,
+        pipe: None = None
+) -> gr.Interface:
+
+    raise NotImplementedError
+
+def create_demo(
+        runner: Runner,
+        pipe: None = None
+) -> gr.TabbedInterface:
+
+    title = "Matcher🎯: Segment Anything with One Shot Using All-Purpose Feature Matching<br> \
+    <div align='center'> \
+    <h2><a href='https://arxiv.org/abs/2305.13310' target='_blank' rel='noopener'>[paper]</a> \
+    <a href='https://github.com/aim-uofa/Matcher' target='_blank' rel='noopener'>[code]</a></h2> \
+    <h2>Matcher can segment anything with one shot by integrating an all-purpose feature extraction model and a class-agnostic segmentation model.</h2> \
+    <br> \
+    </div> \
+    "
+
+    oss_demo = create_oss_demo(runner=runner, pipe=pipe)
+    # vos_demo = create_vos_demo(runner=runner, pipe=pipe)
+    demo = gr.TabbedInterface(
+        [oss_demo,],
+        ['OSS+OPS',], title=title)
+    return demo
+
+
+if __name__ == '__main__':
+    pipe = None
+    HF_TOKEN = os.getenv('HF_TOKEN')
+    runner = Runner(HF_TOKEN)
+    # runner = None
+    demo = create_demo(runner, pipe)
+    demo.launch(enable_queue=False)

oss_ops_inference.py

@@ -0,0 +1,288 @@
+r""" HyperAverageMetercorrelation Squeeze testing code """
+
+import argparse
+import sys
+import os
+from os.path import join
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+
+from segment_anything import SamPredictor, SamAutomaticMaskGenerator
+
+from gradio_demo.Matcher import Matcher
+from matcher.common import utils
+
+import random
+random.seed(0)
+
+
+def default_argument_parser():
+
+    # Arguments parsing
+    parser = argparse.ArgumentParser(description='Matcher Pytorch Implementation for One-shot Segmentation')
+
+    # Dataset parameters
+    parser.add_argument('--datapath', type=str, default='datasets')
+    parser.add_argument('--benchmark', type=str, default='coco',
+                        choices=['fss', 'coco', 'lvis', 'paco_part', 'pascal_part'])
+    parser.add_argument('--bsz', type=int, default=1)
+    parser.add_argument('--nworker', type=int, default=0)
+    parser.add_argument('--fold', type=int, default=0)
+    parser.add_argument('--nshot', type=int, default=1)
+    parser.add_argument('--img-size', type=int, default=518)
+    parser.add_argument('--use_original_imgsize', action='store_true')
+    parser.add_argument('--log-root', type=str, default='output/coco/fold0')
+    parser.add_argument('--visualize', type=int, default=0)
+
+    # DINOv2 and SAM parameters
+    parser.add_argument('--dinov2-weights', type=str, default="models/dinov2_vitl14_pretrain.pth")
+    parser.add_argument('--sam-weights', type=str, default="models/sam_vit_h_4b8939.pth")
+    parser.add_argument('--points_per_side', type=int, default=64)
+    parser.add_argument('--pred_iou_thresh', type=float, default=0.88)
+    parser.add_argument('--sel_stability_score_thresh', type=float, default=0.0)
+    parser.add_argument('--stability_score_thresh', type=float, default=0.95)
+    parser.add_argument('--iou_filter', type=float, default=0.0)
+    parser.add_argument('--box_nms_thresh', type=float, default=1.0)
+    parser.add_argument('--output_layer', type=int, default=3)
+    parser.add_argument('--dense_multimask_output', type=int, default=0)
+    parser.add_argument('--use_dense_mask', type=int, default=0)
+    parser.add_argument('--multimask_output', type=int, default=0)
+
+    # Matcher parameters
+    parser.add_argument('--num_centers', type=int, default=8, help='K centers for kmeans')
+    parser.add_argument('--use_box', action='store_true', help='use box as an extra prompt for sam')
+    parser.add_argument('--use_points_or_centers', action='store_true', help='points:T, center: F')
+    parser.add_argument('--sample-range', type=tuple, default=(4,6), help='sample points number range')
+    parser.add_argument('--max_sample_iterations', type=int, default=30)
+    parser.add_argument('--alpha', type=float, default=1.)
+    parser.add_argument('--beta', type=float, default=0.)
+    parser.add_argument('--exp', type=float, default=0.)
+    parser.add_argument('--emd_filter', type=float, default=0.0, help='use emd_filter')
+    parser.add_argument('--purity_filter', type=float, default=0.0, help='use purity_filter')
+    parser.add_argument('--coverage_filter', type=float, default=0.0, help='use coverage_filter')
+    parser.add_argument('--use_score_filter', action='store_true')
+    parser.add_argument('--deep_score_norm_filter', type=float, default=0.1)
+    parser.add_argument('--deep_score_filter', type=float, default=0.33)
+    parser.add_argument('--topk_scores_threshold', type=float, default=0.7)
+    parser.add_argument('--num_merging_mask', type=int, default=10, help='topk masks for merging')
+
+    args = parser.parse_args()
+    return args
+
+def definite_argument_parser(args, version=1):
+
+    if version==1:
+
+        args.max_sample_iterations = 64
+        args.box_nms_thresh = 0.65
+        args.sample_range = (1, 6)
+        args.topk_scores_threshold = 0.0
+        args.use_dense_mask = 1
+        args.use_points_or_centers = True
+        args.purity_filter = 0.02
+        args.iou_filter = 0.85
+        args.multimask_output = 1
+        args.sel_stability_score_thresh = 0.90
+        args.use_score_filter = True
+        args.alpha = 1.0
+        args.beta = 0.
+        args.exp = 0.
+        args.num_merging_mask = 9
+    elif version == 2:
+        args.max_sample_iterations = 30
+        args.sample_range = (4, 6)
+        args.multimask_output = 0
+        args.alpha = 0.8
+        args.beta = 0.2
+        args.exp = 1.
+        args.num_merging_mask = 10
+    elif version == 3:
+        args.max_sample_iterations = 128
+        args.sample_range = (3, 6)
+        args.use_box = True
+        args.use_points_or_centers = True
+        args.coverage_filter = 0.3
+        args.alpha = 0.5
+        args.beta = 0.5
+        args.exp = 0.
+        args.num_merging_mask = 5
+
+    return args
+
+def preprocess_data(kwargs, args=None):
+
+    img_size = args.img_size
+    transform = transforms.Compose([
+        transforms.Resize(size=(img_size, img_size)),
+        transforms.ToTensor()
+    ])
+
+    support_img = Image.fromarray(kwargs.get("support_img"))
+    query_img_1 = Image.fromarray(kwargs.get("query_img_1"))
+    query_img_2 = Image.fromarray(kwargs.get("query_img_2"))
+
+    support_img_ori_size = (support_img.size[1], support_img.size[0]) # H, W
+    query_img_1_ori_size = (query_img_1.size[1], query_img_1.size[0])
+    query_img_2_ori_size = (query_img_2.size[1], query_img_2.size[0])
+
+
+    support_img = transform(support_img)
+    query_img_1 = transform(query_img_1)
+    query_img_2 = transform(query_img_2)
+
+    support_mask = torch.tensor(kwargs.get("support_mask"))
+    support_mask = F.interpolate(support_mask.unsqueeze(0).float(), support_img.size()[-2:],
+                               mode='nearest') > 0
+    query_imgs = torch.stack([query_img_1, query_img_2], dim=0)
+
+    data = {
+        "support_img": support_img[None, ...],
+        "support_mask": support_mask,
+        "query_imgs": query_imgs,
+        "support_img_ori_size": support_img_ori_size,
+        "query_imgs_ori_size": (query_img_1_ori_size, query_img_2_ori_size),
+    }
+
+    return data
+
+def preprocess_support_mask(data, predictor, version=1):
+
+    if version == 3:
+        return data
+
+    sup_mask = data['support_mask'].squeeze()
+    H, W = sup_mask.shape[-2:]
+    input_points = sup_mask.nonzero().numpy()[:1,::-1]#[:,::-1]
+    input_label = np.array([1]*len(input_points))
+
+    support_img_np = data['support_img'].mul(255).byte()
+    support_img_np = support_img_np.squeeze().permute(1,2,0).cpu().numpy()
+
+    # forward encoder to obtain image feature
+    predictor.reset_image()
+    predictor.set_image(support_img_np)
+
+    # mask, _, _ = predictor.predict(
+    #     point_coords=input_points,
+    #     point_labels=input_label,
+    #     multimask_output=False #True
+    # )
+    mask, _, _ = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_label,
+        multimask_output=True  # True
+    )
+    predictor.reset_image()
+
+    # show_img_point_box_mask(
+    #     support_img_np,
+    #     masks=mask,
+    #     save_path='test1.png',
+    #     mode='mask'
+    # )
+
+    # data['support_mask'] = torch.tensor(mask[:1])[None, ...]
+    data['support_mask'] = torch.tensor(mask[-1:])[None, ...]
+
+    return data
+
+def main_oss_ops(**kwargs):
+
+    args = default_argument_parser()
+    args = definite_argument_parser(args, kwargs.get("version"))
+
+    # Model initialization
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    args.device = device
+
+    # create sam
+    sam = kwargs.get("sam")
+    predictor = SamPredictor(sam)
+    generator = SamAutomaticMaskGenerator(
+        sam,
+        points_per_side=args.points_per_side,
+        points_per_batch=64,
+        pred_iou_thresh=args.pred_iou_thresh,
+        stability_score_thresh=args.stability_score_thresh,
+        stability_score_offset=1.0,
+        sel_stability_score_thresh=args.sel_stability_score_thresh,
+        sel_pred_iou_thresh=args.iou_filter,
+        box_nms_thresh=args.box_nms_thresh,
+        sel_output_layer=args.output_layer,
+        output_layer=args.dense_multimask_output,
+        dense_pred=args.use_dense_mask,
+        multimask_output=args.dense_multimask_output > 0,
+        sel_multimask_output=args.multimask_output > 0,
+    )
+
+    # create dinov2, large
+    dinov2 = kwargs.get("dinov2")
+
+    # create matcher
+    score_filter_cfg = {
+        "emd": args.emd_filter,
+        "purity": args.purity_filter,
+        "coverage": args.coverage_filter,
+        "score_filter": args.use_score_filter,
+        "score": args.deep_score_filter,
+        "score_norm": args.deep_score_norm_filter,
+        "topk_scores_threshold": args.topk_scores_threshold
+    }
+
+    matcher = Matcher(
+        encoder=dinov2,
+        generator=generator,
+        num_centers=args.num_centers,
+        use_box=args.use_box,
+        use_points_or_centers=args.use_points_or_centers,
+        sample_range=args.sample_range,
+        max_sample_iterations=args.max_sample_iterations,
+        alpha=args.alpha,
+        beta=args.beta,
+        exp=args.exp,
+        score_filter_cfg=score_filter_cfg,
+        num_merging_mask=args.num_merging_mask,
+        device=args.device
+    )
+
+    # process data
+    data = preprocess_data(kwargs, args=args)
+    data = preprocess_support_mask(data, predictor, version=kwargs.get("version"))
+
+    # inference
+    with torch.no_grad():
+        utils.fix_randseed(0)
+        pred_masks, pred_mask_lists = [], []
+
+        # support mask
+        support_img_ori_size = data['support_img_ori_size']
+        mask = data['support_mask'].to(predictor.model.device).float()
+        mask = F.interpolate(mask, support_img_ori_size, mode="bilinear", align_corners=False) > 0
+        mask = mask.squeeze(0).cpu().numpy()
+        pred_masks.append(mask)
+        pred_mask_lists.append(None)
+
+        for query_img, query_img_ori_size in zip(data['query_imgs'], data['query_imgs_ori_size']):
+            data['query_img'], data['query_img_ori_size'] = query_img[None, ...], query_img_ori_size
+
+            support_imgs, support_masks = data["support_img"].to(matcher.device)[None, ...], data["support_mask"].to(matcher.device)  # (1, 1, 3, H, W), (1, 1, H, W)
+            query_img, query_img_ori_size = data['query_img'].to(matcher.device), data['query_img_ori_size']  # (1, 3, H, W), img_size
+
+            # 1. Matcher prepare references and target
+            matcher.set_reference(support_imgs, support_masks)
+            matcher.set_target(query_img, query_img_ori_size)
+
+            # 2. Predict mask of target
+            pred_mask, pred_mask_list = matcher.predict()
+            matcher.clear()
+
+            pred_masks.append(pred_mask)
+            pred_mask_lists.append(pred_mask_list)
+
+
+    return pred_masks, pred_mask_lists

runner.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import datetime
+import os
+import pathlib
+import shlex
+import shutil
+import subprocess
+import sys
+
+import gradio as gr
+import slugify
+import torch
+import numpy as np
+import huggingface_hub
+from huggingface_hub import HfApi
+from omegaconf import OmegaConf
+
+from segment_anything import sam_model_registry
+from dinov2.models import vision_transformer as vits
+import dinov2.utils.utils as dinov2_utils
+
+from gradio_demo.oss_ops_inference import main_oss_ops
+
+
+ORIGINAL_SPACE_ID = ''
+SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
+
+
+class Runner:
+    def __init__(self, hf_token: str | None = None):
+        self.hf_token = hf_token
+
+        # self.checkpoint_dir = pathlib.Path('checkpoints')
+        # self.checkpoint_dir.mkdir(exist_ok=True)
+
+        # oss, ops
+        self.prompt_res_g = None
+        self.prompt_mask_g = None
+        self.tar1_res_g = None
+        self.tar2_res_g = None
+        self.version = 1
+
+        self.pred_masks = None
+        self.pred_mask_lists = None
+
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+        sam_checkpoint = "models/sam_vit_h_4b8939.pth"
+        model_type = "default"
+        self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+        self.sam.to(device=device)
+
+        dinov2_kwargs = dict(
+            img_size=518,
+            patch_size=14,
+            init_values=1e-5,
+            ffn_layer='mlp',
+            block_chunks=0,
+            qkv_bias=True,
+            proj_bias=True,
+            ffn_bias=True,
+        )
+        dinov2 = vits.__dict__["vit_large"](**dinov2_kwargs)
+
+        dinov2_utils.load_pretrained_weights(dinov2, "models/dinov2_vitl14_pretrain.pth", "teacher")
+        dinov2.eval()
+        dinov2.to(device=device)
+        self.dinov2 = dinov2
+
+    def inference_oss_ops(self, prompt, target1, target2, version):
+
+        if version == 'version 1 (🔺 multiple instances  🔻 whole, 🔻 part)':
+            self.prompt_res_g, self.tar1_res_g, self.tar2_res_g = prompt['image'], target1, target2
+            self.prompt_mask_g = (prompt['mask'][..., 0] != 0)[None, ...] # 1, H, w
+            self.version = 1
+        elif version == 'version 2 (🔻 multiple instances  🔺 whole, 🔻 part)':
+            self.prompt_res_g, self.tar1_res_g, self.tar2_res_g = prompt['image'], target1, target2
+            self.prompt_mask_g = (prompt['mask'][..., 0] != 0)[None, ...]  # 1, H, w
+            self.version = 2
+        else:
+            self.prompt_res_g, self.tar1_res_g, self.tar2_res_g = prompt['image'], target1, target2
+            self.prompt_mask_g = (prompt['mask'][..., 0] != 0)[None, ...]  # 1, H, w
+            self.version = 3
+
+        self.pred_masks, self.pred_mask_lists = main_oss_ops(
+            sam=self.sam,
+            dinov2=self.dinov2,
+            support_img=self.prompt_res_g,
+            support_mask=self.prompt_mask_g,
+            query_img_1=self.tar1_res_g,
+            query_img_2=self.tar2_res_g,
+            version=self.version
+        )
+
+        text = "Process Successful!"
+
+        return text
+
+
+    def clear_fn(self):
+
+        self.prompt_res_g, self.tar1_res_g, self.tar2_res_g, self.prompt_mask_g = None, None, None, None
+        self.version = 1
+        self.pred_masks = None
+        self.pred_mask_lists = None
+
+        return [None] * 7
+
+
+    def controllable_mask_output(self, k):
+
+        color = np.array([30, 144, 255])
+
+        if self.version != 1:
+
+            prompt_mask_res, tar1_mask_res, tar2_mask_res = self.pred_masks
+
+            h, w = prompt_mask_res.shape[-2:]
+            prompt_mask_img = prompt_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            prompt_mask_res = self.prompt_res_g * 0.5 + prompt_mask_img * 0.5
+
+            h, w = tar1_mask_res.shape[-2:]
+            tar1_mask_img = tar1_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            tar1_mask_res = self.tar1_res_g * 0.5 + tar1_mask_img * 0.5
+
+            h, w = tar2_mask_res.shape[-2:]
+            tar2_mask_img = tar2_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            tar2_mask_res = self.tar2_res_g * 0.5 + tar2_mask_img * 0.5
+
+        else:
+            prompt_mask_res = self.pred_masks[0]
+            tar1_mask_res, tar2_mask_res = self.pred_mask_lists[1:]
+
+            tar1_mask_res = tar1_mask_res[:min(k, len(tar1_mask_res))].sum(0)>0
+            tar2_mask_res = tar2_mask_res[:min(k, len(tar2_mask_res))].sum(0) > 0
+
+            h, w = prompt_mask_res.shape[-2:]
+            prompt_mask_img = prompt_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            prompt_mask_res = self.prompt_res_g * 0.5 + prompt_mask_img * 0.5
+
+            h, w = tar1_mask_res.shape[-2:]
+            tar1_mask_img = tar1_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            tar1_mask_res = self.tar1_res_g * 0.5 + tar1_mask_img * 0.5
+
+            h, w = tar2_mask_res.shape[-2:]
+            tar2_mask_img = tar2_mask_res.reshape(h, w, 1) * color.reshape(1, 1, -1)
+            tar2_mask_res = self.tar2_res_g * 0.5 + tar2_mask_img * 0.5
+
+        return prompt_mask_res/255, tar1_mask_res/255, tar2_mask_res/255
+    
+
+    def inference_vos(self, prompt_vid, vid):
+
+        raise NotImplementedError
+