Upload 37 files

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 regionspot/data/__pycache__/v3det_categories.cpython-38.pyc filter=lfs diff=lfs merge=lfs -text
+images/pic2.png filter=lfs diff=lfs merge=lfs -text
+images/pic3.png filter=lfs diff=lfs merge=lfs -text

configs/Base-RegionSpot.yaml

@@ -0,0 +1,36 @@
+MODEL:
+  META_ARCHITECTURE: "RegionSpot"
+  PIXEL_MEAN: [123.675, 116.280, 103.530]
+  PIXEL_STD: [58.395, 57.120, 57.375]
+  
+SOLVER:
+  IMS_PER_BATCH: 16
+  BASE_LR: 0.000025
+  CHECKPOINT_PERIOD: 50000
+  STEPS: (210000, 250000)
+  MAX_ITER: 270000
+  WARMUP_FACTOR: 0.01
+  WARMUP_ITERS: 1000
+  WEIGHT_DECAY: 0.0001
+  OPTIMIZER: "ADAMW"
+  BACKBONE_MULTIPLIER: 1.0  # keep same with BASE_LR.
+  CLIP_GRADIENTS:
+    ENABLED: True
+    CLIP_TYPE: "full_model"
+    CLIP_VALUE: 1.0
+    NORM_TYPE: 2.0
+SEED: 40244023
+INPUT:
+  MIN_SIZE_TRAIN: (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
+  CROP:
+    ENABLED: False
+    TYPE: "absolute_range"
+    SIZE: (384, 600)
+  FORMAT: "RGB"
+TEST:
+  EVAL_PERIOD: 733000000
+DATALOADER:
+  FILTER_EMPTY_ANNOTATIONS: False
+  NUM_WORKERS: 4
+VERSION: 2
+

configs/eval.yaml

@@ -0,0 +1,21 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  CLIP_TYPE: CLIP_400M_Large
+  TRAINING: False
+  BOX_TYPE: 'PRED_BOX'
+  MASK_ON: True
+DATASETS:  # LVIS
+  TRAIN: ("lvis_v1_train",)
+  TEST: ("lvis_v1_val",)
+DATALOADER:
+  SAMPLER_TRAIN: "RepeatFactorTrainingSampler"
+  REPEAT_THRESHOLD: 0.001
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+TEST:  # LVIS
+  EVAL_PERIOD: 0  # disable eval during train since long time
+
+OUTPUT_DIR: './output/eval'
+

configs/objects365_bb.yaml

@@ -0,0 +1,18 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  CLIP_TYPE: CLIP_400M
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_obj365_bb'

configs/objects365_bl.yaml

@@ -0,0 +1,18 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  CLIP_TYPE: CLIP_400M_Large
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_obj365_bl'

configs/objects365_bl_336.yaml

@@ -0,0 +1,19 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  CLIP_TYPE: CLIP_400M_Large_336
+  CLIP_INPUT_SIZE: 336
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_obj365_bl_336'

configs/objects365_v3det_openimages_bb.yaml

@@ -0,0 +1,19 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  WEIGHTS: "./output/regionspot_obj365_bb/model_final.pth"
+  CLIP_TYPE: CLIP_400M
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train", "v3det_train","openimages_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_alldata_bb'

configs/objects365_v3det_openimages_bl.yaml

@@ -0,0 +1,19 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  WEIGHTS: "./output/regionspot_obj365_bl/model_final.pth"
+  CLIP_TYPE: CLIP_400M_Large
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train", "v3det_train","openimages_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_alldata_bl'

configs/objects365_v3det_openimages_bl_336.yaml

@@ -0,0 +1,20 @@
+_BASE_: "Base-RegionSpot.yaml"
+MODEL:
+  WEIGHTS: "./output/regionspot_obj365_bl_336/model_final.pth"
+  CLIP_TYPE: CLIP_400M_Large_336
+  CLIP_INPUT_SIZE: 336
+DATALOADER:
+  SAMPLER_TRAIN: "MultiDatasetSampler"
+DATASETS:
+  TRAIN: ("objects365_train", "v3det_train","openimages_train",)
+  TEST:  ()
+TEST:
+  EVAL_PERIOD: 0
+SOLVER:
+  STEPS: (350000, 420000)
+  MAX_ITER: 450000  
+INPUT:
+  CROP:
+    ENABLED: True
+  FORMAT: "RGB"
+OUTPUT_DIR: './output/regionspot_alldata_bl_336'

demo.py

@@ -0,0 +1,67 @@
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+from regionspot.modeling.regionspot import build_regionspot_model
+from regionspot import RegionSpot_Predictor
+# Function to show masks on an image
+def show_mask(mask, ax, random_color=False):
+    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)
+
+# Function to show points on an 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)
+
+# Function to show bounding boxes on an image
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - x0, box[3] - y0
+    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor='none', lw=2))
+
+# Read image and set up model
+image = cv2.imread('assets/image.jpg')
+image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # Convert image to RGB format
+# Multiple boxes
+box_prompt = np.array([[64, 926, 804, 1978], [1237, 490, 1615, 771.], [1510, 64, 1670, 167]]) 
+ckpt_path = '/path/to/model_weights.pth'
+clip_type = 'CLIP_400M_Large_336'
+clip_input_size = 336
+custom_vocabulary =  ["Smoothie bowl", "Banana", "Strawberry", "Chia seeds", "Shredded coconut", "Wooden spoons", "Grapefruit", "Goji berries", "Flaxseeds seeds"]
+
+# Build and initialize the model
+model, msg = build_regionspot_model(is_training=False, image_size=clip_input_size, clip_type=clip_type, pretrain_ckpt=ckpt_path, custom_vocabulary=custom_vocabulary)
+
+# Create predictor and set image
+predictor = RegionSpot_Predictor(model.cuda())
+predictor.set_image(image, clip_input_size=clip_input_size)
+
+# Prediction based on box prompt
+masks, mask_iou_score, class_score, class_index = predictor.predict(
+    point_coords=None,
+    point_labels=None,
+    box=box_prompt,
+    multimask_output=False,  
+)
+# Extract class name and display image with masks and box
+fig, ax = plt.subplots(figsize=(10, 10))
+ax.imshow(image)
+for idx in range(len(box_prompt)):
+    show_mask(masks[idx], ax)  
+    show_box(box_prompt[idx], ax)  # Assuming box_prompt contains all your boxes
+    # You might want to modify the text display for multiple classes as well
+    class_name = custom_vocabulary[int(class_index[idx])]
+    ax.text(box_prompt[idx][0], box_prompt[idx][1] - 10, class_name, color='white', fontsize=14, bbox=dict(facecolor='green', edgecolor='green', alpha=0.6))
+
+ax.axis('off')
+plt.show()
+fig.savefig('result.png')
+plt.close(fig)
+

flagged/log.csv

@@ -0,0 +1,2 @@
+input_img,output,flag,username,timestamp
+"{""path"":""flagged\\input_img\\0fce6cef302bd30bbf85\\000000052891.jpg"",""url"":""http://127.0.0.1:7862/file=C:\\Users\\10051\\AppData\\Local\\Temp\\gradio\\f11efb7c9dd110c5865f9afd384282cda8ca1f25\\000000052891.jpg"",""size"":153377,""orig_name"":""000000052891.jpg"",""mime_type"":""""}","{""path"":""flagged\\output\\440681fbd96cfcf67ac5\\image.png"",""url"":null,""size"":null,""orig_name"":null,""mime_type"":null}",,,2023-11-10 12:08:05.538914

tools/re_save_ckpt.py

@@ -0,0 +1,22 @@
+import torch
+
+pretrain_ckpt = './pretrained_model/model_final.pth'
+checkpoint = torch.load(pretrain_ckpt, map_location='cpu')
+
+# Remove specific keys from the top-level dictionary
+top_level_keys_to_remove = ['trainer', 'iteration']
+for key in top_level_keys_to_remove:
+    if key in checkpoint:
+        del checkpoint[key]
+
+# Remove keys that start with 'clip_model' and 'sam' from the checkpoint's 'model' dictionary
+model_keys_to_remove = ['model.clip_model', 'model.sam']
+for key in list(checkpoint['model'].keys()):  # Use list to copy keys
+    if any(key.startswith(to_remove) for to_remove in model_keys_to_remove):
+        print(key)
+        del checkpoint['model'][key]
+
+# Save the modified checkpoint back to a file
+modified_ckpt_path = './pretrained_model/model_final_modified.pth'
+torch.save(checkpoint, modified_ckpt_path)
+print(checkpoint['model'].keys())

train_net.py

@@ -0,0 +1,297 @@
+"""
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+This script is a simplified version of the training script in detectron2/tools.
+"""
+import os
+import itertools
+import weakref
+from typing import Any, Dict, List, Set
+import logging
+from collections import OrderedDict
+
+import torch
+from fvcore.nn.precise_bn import get_bn_modules
+
+import detectron2.utils.comm as comm
+from detectron2.utils.logger import setup_logger
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import get_cfg
+from detectron2.data import build_detection_train_loader
+from regionspot import build_custom_train_loader
+
+from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, launch, create_ddp_model, \
+    AMPTrainer, SimpleTrainer, hooks
+from detectron2.evaluation import COCOEvaluator, LVISEvaluator, verify_results
+from detectron2.solver.build import maybe_add_gradient_clipping
+from detectron2.modeling import build_model
+from regionspot.data import objects365
+from regionspot.data import openimages
+from regionspot.data import v3det
+
+
+from regionspot import RegionSpotDatasetMapper, add_regionspot_config, RegionSpotWithTTA
+from regionspot.util.model_ema import add_model_ema_configs, may_build_model_ema, may_get_ema_checkpointer, EMAHook, \
+    apply_model_ema_and_restore, EMADetectionCheckpointer
+
+
+class Trainer(DefaultTrainer):
+    """ Extension of the Trainer class adapted to RegionSpot. """
+
+    def __init__(self, cfg):
+        """
+        Args:
+            cfg (CfgNode):
+        """
+        super(DefaultTrainer, self).__init__()  # call grandfather's `__init__` while avoid father's `__init()`
+        logger = logging.getLogger("detectron2")
+        if not logger.isEnabledFor(logging.INFO):  # setup_logger is not called for d2
+            setup_logger()
+        cfg = DefaultTrainer.auto_scale_workers(cfg, comm.get_world_size())
+        # Assume these objects must be constructed in this order.
+        model = self.build_model(cfg)
+        optimizer = self.build_optimizer(cfg, model)
+        data_loader = self.build_train_loader(cfg)
+
+        model = create_ddp_model(model, broadcast_buffers=False)
+        self._trainer = (AMPTrainer if cfg.SOLVER.AMP.ENABLED else SimpleTrainer)(
+            model, data_loader, optimizer
+        )
+
+        self.scheduler = self.build_lr_scheduler(cfg, optimizer)
+
+        ########## EMA ############
+        kwargs = {
+            'trainer': weakref.proxy(self),
+        }
+        kwargs.update(may_get_ema_checkpointer(cfg, model))
+        self.checkpointer = DetectionCheckpointer(
+            # Assume you want to save checkpoints together with logs/statistics
+            model,
+            cfg.OUTPUT_DIR,
+            **kwargs,
+            # trainer=weakref.proxy(self),
+        )
+        self.start_iter = 0
+        self.max_iter = cfg.SOLVER.MAX_ITER
+        self.cfg = cfg
+
+        self.register_hooks(self.build_hooks())
+
+    @classmethod
+    def build_model(cls, cfg):
+        """
+        Returns:
+            torch.nn.Module:
+
+        It now calls :func:`detectron2.modeling.build_model`.
+        Overwrite it if you'd like a different model.
+        """
+        model = build_model(cfg)
+        logger = logging.getLogger(__name__)
+        logger.info("Model:\n{}".format(model))
+        # setup EMA
+        may_build_model_ema(cfg, model)
+        return model
+
+    @classmethod
+    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
+        """
+        Create evaluator(s) for a given dataset.
+        This uses the special metadata "evaluator_type" associated with each builtin dataset.
+        For your own dataset, you can simply create an evaluator manually in your
+        script and do not have to worry about the hacky if-else logic here.
+        """
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+        if 'lvis' in dataset_name:
+            return LVISEvaluator(dataset_name, cfg, True, output_folder)
+        else:
+            return COCOEvaluator(dataset_name, cfg, True, output_folder)
+
+    @classmethod
+    def build_train_loader(cls, cfg): 
+        mapper = RegionSpotDatasetMapper(cfg, is_train=True)
+        if cfg.DATALOADER.SAMPLER_TRAIN in ['TrainingSampler', 'RepeatFactorTrainingSampler']:
+            data_loader = build_detection_train_loader(cfg, mapper=mapper)
+        else:
+            data_loader = build_custom_train_loader(cfg, mapper=mapper)
+        return data_loader
+    @classmethod
+    def build_optimizer(cls, cfg, model):
+        params: List[Dict[str, Any]] = []
+        memo: Set[torch.nn.parameter.Parameter] = set()
+        for key, value in model.named_parameters(recurse=True):
+            if not value.requires_grad:
+                continue
+            # Avoid duplicating parameters
+            if value in memo:
+                continue
+            memo.add(value)
+            lr = cfg.SOLVER.BASE_LR
+            weight_decay = cfg.SOLVER.WEIGHT_DECAY
+            if "backbone" in key:
+                lr = lr * cfg.SOLVER.BACKBONE_MULTIPLIER
+            params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
+
+        def maybe_add_full_model_gradient_clipping(optim):  # optim: the optimizer class
+            # detectron2 doesn't have full model gradient clipping now
+            clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE
+            enable = (
+                    cfg.SOLVER.CLIP_GRADIENTS.ENABLED
+                    and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model"
+                    and clip_norm_val > 0.0
+            )
+
+            class FullModelGradientClippingOptimizer(optim):
+                def step(self, closure=None):
+                    all_params = itertools.chain(*[x["params"] for x in self.param_groups])
+                    torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
+                    super().step(closure=closure)
+
+            return FullModelGradientClippingOptimizer if enable else optim
+
+        optimizer_type = cfg.SOLVER.OPTIMIZER
+        if optimizer_type == "SGD":
+            optimizer = maybe_add_full_model_gradient_clipping(torch.optim.SGD)(
+                params, cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM
+            )
+        elif optimizer_type == "ADAMW":
+            optimizer = maybe_add_full_model_gradient_clipping(torch.optim.AdamW)(
+                params, cfg.SOLVER.BASE_LR
+            )
+        else:
+            raise NotImplementedError(f"no optimizer type {optimizer_type}")
+        if not cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model":
+            optimizer = maybe_add_gradient_clipping(cfg, optimizer)
+        return optimizer
+
+    @classmethod
+    def ema_test(cls, cfg, model, evaluators=None):
+        # model with ema weights
+        logger = logging.getLogger("detectron2.trainer")
+        if cfg.MODEL_EMA.ENABLED:
+            logger.info("Run evaluation with EMA.")
+            with apply_model_ema_and_restore(model):
+                results = cls.test(cfg, model, evaluators=evaluators)
+        else:
+            results = cls.test(cfg, model, evaluators=evaluators)
+        return results
+
+    @classmethod
+    def test_with_TTA(cls, cfg, model):
+        logger = logging.getLogger("detectron2.trainer")
+        logger.info("Running inference with test-time augmentation ...")
+        model = RegionSpotWithTTA(cfg, model)
+        evaluators = [
+            cls.build_evaluator(
+                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
+            )
+            for name in cfg.DATASETS.TEST
+        ]
+        if cfg.MODEL_EMA.ENABLED:
+            cls.ema_test(cfg, model, evaluators)
+        else:
+            res = cls.test(cfg, model, evaluators)
+        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
+        return res
+
+    def build_hooks(self):
+        """
+        Build a list of default hooks, including timing, evaluation,
+        checkpointing, lr scheduling, precise BN, writing events.
+
+        Returns:
+            list[HookBase]:
+        """
+        cfg = self.cfg.clone()
+        cfg.defrost()
+        cfg.DATALOADER.NUM_WORKERS = 0  # save some memory and time for PreciseBN
+
+        ret = [
+            hooks.IterationTimer(),
+            EMAHook(self.cfg, self.model) if cfg.MODEL_EMA.ENABLED else None,  # EMA hook
+            hooks.LRScheduler(),
+            hooks.PreciseBN(
+                # Run at the same freq as (but before) evaluation.
+                cfg.TEST.EVAL_PERIOD,
+                self.model,
+                # Build a new data loader to not affect training
+                self.build_train_loader(cfg),
+                cfg.TEST.PRECISE_BN.NUM_ITER,
+            )
+            if cfg.TEST.PRECISE_BN.ENABLED and get_bn_modules(self.model)
+            else None,
+        ]
+
+        # Do PreciseBN before checkpointer, because it updates the model and need to
+        # be saved by checkpointer.
+        # This is not always the best: if checkpointing has a different frequency,
+        # some checkpoints may have more precise statistics than others.
+        if comm.is_main_process():
+            ret.append(hooks.PeriodicCheckpointer(self.checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD))
+
+        def test_and_save_results():
+            self._last_eval_results = self.test(self.cfg, self.model)
+            return self._last_eval_results
+
+        # Do evaluation after checkpointer, because then if it fails,
+        # we can use the saved checkpoint to debug.
+        ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
+
+        if comm.is_main_process():
+            # Here the default print/log frequency of each writer is used.
+            # run writers in the end, so that evaluation metrics are written
+            ret.append(hooks.PeriodicWriter(self.build_writers(), period=20))
+        return ret
+
+
+def setup(args):
+    """
+    Create configs and perform basic setups.
+    """
+    cfg = get_cfg()
+    add_regionspot_config(cfg)
+    add_model_ema_configs(cfg)
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    default_setup(cfg, args)
+    return cfg
+
+
+def main(args):
+    cfg = setup(args)
+
+    if args.eval_only:
+        model = Trainer.build_model(cfg)
+        kwargs = may_get_ema_checkpointer(cfg, model)
+        if cfg.MODEL_EMA.ENABLED:
+            EMADetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR, **kwargs).resume_or_load(cfg.MODEL.WEIGHTS,
+                                                                                              resume=args.resume)
+        else:
+            DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR, **kwargs).resume_or_load(cfg.MODEL.WEIGHTS,
+                                                                                           resume=args.resume)
+        res = Trainer.ema_test(cfg, model)
+        if cfg.TEST.AUG.ENABLED:
+            res.update(Trainer.test_with_TTA(cfg, model))
+        if comm.is_main_process():
+            verify_results(cfg, res)
+        return res
+
+    trainer = Trainer(cfg)
+    trainer.resume_or_load(resume=args.resume)
+    return trainer.train()
+
+
+if __name__ == "__main__":
+    args = default_argument_parser().parse_args()
+    print("Command Line Args:", args)
+    launch(
+        main,
+        args.num_gpus,
+        num_machines=args.num_machines,
+        machine_rank=args.machine_rank,
+        dist_url=args.dist_url,
+        args=(args,),
+    )

utils/tools.py

@@ -0,0 +1,442 @@
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import cv2
+import torch
+import os
+import sys
+import clip
+
+
+def convert_box_xywh_to_xyxy(box):
+    if len(box) == 4:
+        return [box[0], box[1], box[0] + box[2], box[1] + box[3]]
+    else:
+        result = []
+        for b in box:
+            b = convert_box_xywh_to_xyxy(b)
+            result.append(b)               
+    return result
+
+
+def segment_image(image, bbox):
+    image_array = np.array(image)
+    segmented_image_array = np.zeros_like(image_array)
+    x1, y1, x2, y2 = bbox
+    segmented_image_array[y1:y2, x1:x2] = image_array[y1:y2, x1:x2]
+    segmented_image = Image.fromarray(segmented_image_array)
+    black_image = Image.new("RGB", image.size, (255, 255, 255))
+    # transparency_mask = np.zeros_like((), dtype=np.uint8)
+    transparency_mask = np.zeros(
+        (image_array.shape[0], image_array.shape[1]), dtype=np.uint8
+    )
+    transparency_mask[y1:y2, x1:x2] = 255
+    transparency_mask_image = Image.fromarray(transparency_mask, mode="L")
+    black_image.paste(segmented_image, mask=transparency_mask_image)
+    return black_image
+
+
+def format_results(result, filter=0):
+    annotations = []
+    n = len(result.masks.data)
+    for i in range(n):
+        annotation = {}
+        mask = result.masks.data[i] == 1.0
+
+        if torch.sum(mask) < filter:
+            continue
+        annotation["id"] = i
+        annotation["segmentation"] = mask.cpu().numpy()
+        annotation["bbox"] = result.boxes.data[i]
+        annotation["score"] = result.boxes.conf[i]
+        annotation["area"] = annotation["segmentation"].sum()
+        annotations.append(annotation)
+    return annotations
+
+
+def filter_masks(annotations):  # filter the overlap mask
+    annotations.sort(key=lambda x: x["area"], reverse=True)
+    to_remove = set()
+    for i in range(0, len(annotations)):
+        a = annotations[i]
+        for j in range(i + 1, len(annotations)):
+            b = annotations[j]
+            if i != j and j not in to_remove:
+                # check if
+                if b["area"] < a["area"]:
+                    if (a["segmentation"] & b["segmentation"]).sum() / b[
+                        "segmentation"
+                    ].sum() > 0.8:
+                        to_remove.add(j)
+
+    return [a for i, a in enumerate(annotations) if i not in to_remove], to_remove
+
+
+def get_bbox_from_mask(mask):
+    mask = mask.astype(np.uint8)
+    contours, hierarchy = cv2.findContours(
+        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
+    x1, y1, w, h = cv2.boundingRect(contours[0])
+    x2, y2 = x1 + w, y1 + h
+    if len(contours) > 1:
+        for b in contours:
+            x_t, y_t, w_t, h_t = cv2.boundingRect(b)
+            # 将多个bbox合并成一个
+            x1 = min(x1, x_t)
+            y1 = min(y1, y_t)
+            x2 = max(x2, x_t + w_t)
+            y2 = max(y2, y_t + h_t)
+        h = y2 - y1
+        w = x2 - x1
+    return [x1, y1, x2, y2]
+
+
+def fast_process(
+    annotations, args, mask_random_color, bbox=None, points=None, edges=False
+):
+    if isinstance(annotations[0], dict):
+        annotations = [annotation["segmentation"] for annotation in annotations]
+    result_name = os.path.basename(args.img_path)
+    image = cv2.imread(args.img_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    original_h = image.shape[0]
+    original_w = image.shape[1]
+    if sys.platform == "darwin":
+            plt.switch_backend("TkAgg")
+    plt.figure(figsize=(original_w/100, original_h/100))
+    # Add subplot with no margin.
+    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+    plt.margins(0, 0)
+    plt.gca().xaxis.set_major_locator(plt.NullLocator())
+    plt.gca().yaxis.set_major_locator(plt.NullLocator())
+    plt.imshow(image)
+    if args.better_quality == True:
+        if isinstance(annotations[0], torch.Tensor):
+            annotations = np.array(annotations.cpu())
+        for i, mask in enumerate(annotations):
+            mask = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
+            )
+            annotations[i] = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
+            )
+    if args.device == "cpu":
+        annotations = np.array(annotations)
+        fast_show_mask(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            points=points,
+            point_label=args.point_label,
+            retinamask=args.retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    else:
+        if isinstance(annotations[0], np.ndarray):
+            annotations = torch.from_numpy(annotations)
+        fast_show_mask_gpu(
+            annotations,
+            plt.gca(),
+            random_color=args.randomcolor,
+            bbox=bbox,
+            points=points,
+            point_label=args.point_label,
+            retinamask=args.retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    if isinstance(annotations, torch.Tensor):
+        annotations = annotations.cpu().numpy()
+    if args.withContours == True:
+        contour_all = []
+        temp = np.zeros((original_h, original_w, 1))
+        for i, mask in enumerate(annotations):
+            if type(mask) == dict:
+                mask = mask["segmentation"]
+            annotation = mask.astype(np.uint8)
+            if args.retina == False:
+                annotation = cv2.resize(
+                    annotation,
+                    (original_w, original_h),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            contours, hierarchy = cv2.findContours(
+                annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+            )
+            for contour in contours:
+                contour_all.append(contour)
+        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
+        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
+        contour_mask = temp / 255 * color.reshape(1, 1, -1)
+        plt.imshow(contour_mask)
+
+    save_path = args.output
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+    plt.axis("off")
+    fig = plt.gcf()
+    plt.draw()
+    
+    try:
+        buf = fig.canvas.tostring_rgb()
+    except AttributeError:
+        fig.canvas.draw()
+        buf = fig.canvas.tostring_rgb()
+    
+    cols, rows = fig.canvas.get_width_height()
+    img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
+    cv2.imwrite(os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))
+
+
+# CPU post process
+def fast_show_mask(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    points=None,
+    point_label=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    msak_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    # 将annotation 按照面积 排序
+    areas = np.sum(annotation, axis=(1, 2))
+    sorted_indices = np.argsort(areas)
+    annotation = annotation[sorted_indices]
+
+    index = (annotation != 0).argmax(axis=0)
+    if random_color == True:
+        color = np.random.random((msak_sum, 1, 1, 3))
+    else:
+        color = np.ones((msak_sum, 1, 1, 3)) * np.array(
+            [30 / 255, 144 / 255, 255 / 255]
+        )
+    transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
+    visual = np.concatenate([color, transparency], axis=-1)
+    mask_image = np.expand_dims(annotation, -1) * visual
+
+    show = np.zeros((height, weight, 4))
+    h_indices, w_indices = np.meshgrid(
+        np.arange(height), np.arange(weight), indexing="ij"
+    )
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    show[h_indices, w_indices, :] = mask_image[indices]
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+    # draw point
+    if points is not None:
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
+            s=20,
+            c="y",
+        )
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
+            s=20,
+            c="m",
+        )
+
+    if retinamask == False:
+        show = cv2.resize(
+            show, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    ax.imshow(show)
+
+
+def fast_show_mask_gpu(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    points=None,
+    point_label=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    msak_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    areas = torch.sum(annotation, dim=(1, 2))
+    sorted_indices = torch.argsort(areas, descending=False)
+    annotation = annotation[sorted_indices]
+    # 找每个位置第一个非零值下标
+    index = (annotation != 0).to(torch.long).argmax(dim=0)
+    if random_color == True:
+        color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
+    else:
+        color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor(
+            [30 / 255, 144 / 255, 255 / 255]
+        ).to(annotation.device)
+    transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
+    visual = torch.cat([color, transparency], dim=-1)
+    mask_image = torch.unsqueeze(annotation, -1) * visual
+    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
+    show = torch.zeros((height, weight, 4)).to(annotation.device)
+    h_indices, w_indices = torch.meshgrid(
+        torch.arange(height), torch.arange(weight), indexing="ij"
+    )
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    show[h_indices, w_indices, :] = mask_image[indices]
+    show_cpu = show.cpu().numpy()
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+    # draw point
+    if points is not None:
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
+            s=20,
+            c="y",
+        )
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
+            s=20,
+            c="m",
+        )
+    if retinamask == False:
+        show_cpu = cv2.resize(
+            show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    ax.imshow(show_cpu)
+
+
+# clip
+@torch.no_grad()
+def retriev(
+    model, preprocess, elements: [Image.Image], search_text: str, device
+):
+    preprocessed_images = [preprocess(image).to(device) for image in elements]
+    tokenized_text = clip.tokenize([search_text]).to(device)
+    stacked_images = torch.stack(preprocessed_images)
+    image_features = model.encode_image(stacked_images)
+    text_features = model.encode_text(tokenized_text)
+    image_features /= image_features.norm(dim=-1, keepdim=True)
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    probs = 100.0 * image_features @ text_features.T
+    return probs[:, 0].softmax(dim=0)
+
+
+def crop_image(annotations, image_like):
+    if isinstance(image_like, str):
+        image = Image.open(image_like)
+    else:
+        image = image_like
+    ori_w, ori_h = image.size
+    mask_h, mask_w = annotations[0]["segmentation"].shape
+    if ori_w != mask_w or ori_h != mask_h:
+        image = image.resize((mask_w, mask_h))
+    cropped_boxes = []
+    cropped_images = []
+    not_crop = []
+    origin_id = []
+    for _, mask in enumerate(annotations):
+        if np.sum(mask["segmentation"]) <= 100:
+            continue
+        origin_id.append(_)
+        bbox = get_bbox_from_mask(mask["segmentation"])  # mask 的 bbox
+        cropped_boxes.append(segment_image(image, bbox))  # 保存裁剪的图片
+        # cropped_boxes.append(segment_image(image,mask["segmentation"]))
+        cropped_images.append(bbox)  # 保存裁剪的图片的bbox
+    return cropped_boxes, cropped_images, not_crop, origin_id, annotations
+
+
+def box_prompt(masks, bbox, target_height, target_width):
+    h = masks.shape[1]
+    w = masks.shape[2]
+    if h != target_height or w != target_width:
+        bbox = [
+            int(bbox[0] * w / target_width),
+            int(bbox[1] * h / target_height),
+            int(bbox[2] * w / target_width),
+            int(bbox[3] * h / target_height),
+        ]
+    bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
+    bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
+    bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
+    bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h
+
+    # IoUs = torch.zeros(len(masks), dtype=torch.float32)
+    bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+
+    masks_area = torch.sum(masks[:, bbox[1] : bbox[3], bbox[0] : bbox[2]], dim=(1, 2))
+    orig_masks_area = torch.sum(masks, dim=(1, 2))
+
+    union = bbox_area + orig_masks_area - masks_area
+    IoUs = masks_area / union
+    max_iou_index = torch.argmax(IoUs)
+
+    return masks[max_iou_index].cpu().numpy(), max_iou_index
+
+
+def point_prompt(masks, points, point_label, target_height, target_width):  # numpy 处理
+    h = masks[0]["segmentation"].shape[0]
+    w = masks[0]["segmentation"].shape[1]
+    if h != target_height or w != target_width:
+        points = [
+            [int(point[0] * w / target_width), int(point[1] * h / target_height)]
+            for point in points
+        ]
+    onemask = np.zeros((h, w))
+    masks = sorted(masks, key=lambda x: x['area'], reverse=True)
+    for i, annotation in enumerate(masks):
+        if type(annotation) == dict:
+            mask = annotation['segmentation']
+        else:
+            mask = annotation
+        for i, point in enumerate(points):
+            if mask[point[1], point[0]] == 1 and point_label[i] == 1:
+                onemask[mask] = 1
+            if mask[point[1], point[0]] == 1 and point_label[i] == 0:
+                onemask[mask] = 0
+    onemask = onemask >= 1
+    return onemask, 0
+
+
+def text_prompt(annotations, text, img_path, device, wider=False, threshold=0.9):
+    cropped_boxes, cropped_images, not_crop, origin_id, annotations_ = crop_image(
+        annotations, img_path
+    )
+    clip_model, preprocess = clip.load("./weights/CLIP_ViT_B_32.pt", device=device)
+    scores = retriev(
+        clip_model, preprocess, cropped_boxes, text, device=device
+    )
+    max_idx = scores.argsort()
+    max_idx = max_idx[-1]
+    max_idx = origin_id[int(max_idx)]
+
+    # find the biggest mask which contains the mask with max score
+    if wider:
+        mask0 = annotations_[max_idx]["segmentation"]
+        area0 = np.sum(mask0)
+        areas = [(i, np.sum(mask["segmentation"])) for i, mask in enumerate(annotations_) if i in origin_id]
+        areas = sorted(areas, key=lambda area: area[1], reverse=True)
+        indices = [area[0] for area in areas]
+        for index in indices:
+            if index == max_idx or np.sum(annotations_[index]["segmentation"] & mask0) / area0 > threshold:
+                max_idx = index
+                break
+
+    return annotations_[max_idx]["segmentation"], max_idx

utils/tools_gradio.py

@@ -0,0 +1,203 @@
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import cv2
+import torch
+
+fdic = {
+    "family" : "Impact",
+    "style" : "italic",
+    "size" : 15,
+    "color" : "yellow",
+    "weight" : "bold"
+}
+
+def fast_process(
+    annotations,
+    image,
+    device,
+    scale,
+    better_quality=False,
+    mask_random_color=True,
+    bbox=None,
+    use_retina=True,
+    withContours=True,
+    label = None,
+):
+    if isinstance(annotations[0], dict):
+        annotations = [annotation['segmentation'] for annotation in annotations]
+
+    original_h = image.height
+    original_w = image.width
+    if better_quality:
+        if isinstance(annotations[0], torch.Tensor):
+            annotations = np.array(annotations.cpu())
+        for i, mask in enumerate(annotations):
+            mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
+            annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
+    if device == 'cpu':
+        annotations = np.array(annotations)
+        inner_mask = fast_show_mask(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            retinamask=use_retina,
+            target_height=original_h,
+            target_width=original_w,
+            label = label,
+        )
+    else:
+        if isinstance(annotations[0], np.ndarray):
+            annotations = torch.from_numpy(annotations)
+        inner_mask, figure = fast_show_mask_gpu(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            retinamask=use_retina,
+            target_height=original_h,
+            target_width=original_w,
+            label = label,
+        )
+    if isinstance(annotations, torch.Tensor):
+        annotations = annotations.cpu().numpy()
+
+    if withContours:
+        contour_all = []
+        temp = np.zeros((original_h, original_w, 1))
+        for i, mask in enumerate(annotations):
+            if type(mask) == dict:
+                mask = mask['segmentation']
+            annotation = mask.astype(np.uint8)
+            if use_retina == False:
+                annotation = cv2.resize(
+                    annotation,
+                    (original_w, original_h),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            contours, _ = cv2.findContours(annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+            for contour in contours:
+                contour_all.append(contour)
+        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2 // scale)
+        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.9])
+        contour_mask = temp / 255 * color.reshape(1, 1, -1)
+
+    image = image.convert('RGBA')
+    overlay_inner = Image.fromarray((inner_mask * 255).astype(np.uint8), 'RGBA')
+    image.paste(overlay_inner, (0, 0), overlay_inner)
+
+    if withContours:
+        overlay_contour = Image.fromarray((contour_mask * 255).astype(np.uint8), 'RGBA')
+        image.paste(overlay_contour, (0, 0), overlay_contour)
+
+    plt.figure(figsize=(16, 10))
+    plt.imshow(image)
+    ax = plt.gca()
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
+        ax.text(x1, y1, f"{label}", fontdict=fdic)
+    pic = plt.gcf()
+    pic.canvas.draw()
+    w,h = pic.canvas.get_width_height()
+    image = Image.frombytes('RGB', (w,h), pic.canvas.tostring_rgb())
+
+    return image
+
+
+# CPU post process
+def fast_show_mask(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+    label = None,
+):
+    mask_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    # 将annotation 按照面积 排序
+    areas = np.sum(annotation, axis=(1, 2))
+    sorted_indices = np.argsort(areas)[::1]
+    annotation = annotation[sorted_indices]
+
+    index = (annotation != 0).argmax(axis=0)
+    if random_color:
+        color = np.random.random((mask_sum, 1, 1, 3))
+    else:
+        color = np.ones((mask_sum, 1, 1, 3)) * np.array([30 / 255, 144 / 255, 255 / 255])
+    transparency = np.ones((mask_sum, 1, 1, 1)) * 0.6
+    visual = np.concatenate([color, transparency], axis=-1)
+    mask_image = np.expand_dims(annotation, -1) * visual
+
+    mask = np.zeros((height, weight, 4))
+
+    h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+
+    mask[h_indices, w_indices, :] = mask_image[indices]
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
+        # ax.text(x, y, f"{label}: {round(prediction['score']*100, 1)}%", fontdict=fdic)
+        ax.text(x, y, f"{label}", fontdict=fdic)
+
+    if not retinamask:
+        mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
+
+    return mask
+
+
+def fast_show_mask_gpu(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+    label = None,
+):
+    device = annotation.device
+    mask_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    areas = torch.sum(annotation, dim=(1, 2))
+    sorted_indices = torch.argsort(areas, descending=False)
+    annotation = annotation[sorted_indices]
+    # 找每个位置第一个非零值下标
+    index = (annotation != 0).to(torch.long).argmax(dim=0)
+    if random_color:
+        color = torch.rand((mask_sum, 1, 1, 3)).to(device)
+    else:
+        color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
+            [30 / 255, 144 / 255, 255 / 255]
+        ).to(device)
+    transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
+    visual = torch.cat([color, transparency], dim=-1)
+    mask_image = torch.unsqueeze(annotation, -1) * visual
+    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
+    mask = torch.zeros((height, weight, 4)).to(device)
+    h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    mask[h_indices, w_indices, :] = mask_image[indices]
+    mask_cpu = mask.cpu().numpy()
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=3
+            )
+        )
+        ax.text(x1, y1, f"{label}", fontdict=fdic)
+        # ax.text(x1, y1, f"{label}")
+    if not retinamask:
+        mask_cpu = cv2.resize(
+            mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    return mask_cpu, ax