Add code

app.py

@@ -0,0 +1,252 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Gradio application."""
+
+import argparse
+import multiprocessing as mp
+import os
+import time
+
+import numpy as np
+import torch
+
+from tokenize_anything import test_engine
+from tokenize_anything.utils.image import im_rescale
+from tokenize_anything.utils.image import im_vstack
+
+
+def parse_args():
+    """Parse arguments."""
+    parser = argparse.ArgumentParser(description="Launch gradio app.")
+    parser.add_argument("--model-type", type=str, default="tap_vit_l")
+    parser.add_argument("--checkpoint", type=str, default="models/tap_vit_l_03f8ec.pkl")
+    parser.add_argument("--concept", type=str, default="concepts/merged_2560.pkl")
+    parser.add_argument("--device", nargs="+", type=int, default=[0], help="Index of devices.")
+    return parser.parse_args()
+
+
+class Predictor(object):
+    """Predictor."""
+
+    def __init__(self, model, kwargs):
+        self.model = model
+        self.kwargs = kwargs
+        self.batch_size = kwargs.get("batch_size", 256)
+        self.model.concept_projector.reset_weights(kwargs["concept_weights"])
+        self.model.text_decoder.reset_cache(max_batch_size=self.batch_size)
+
+    def preprocess_images(self, imgs):
+        """Preprocess the inference images."""
+        im_batch, im_shapes, im_scales = [], [], []
+        for img in imgs:
+            scaled_imgs, scales = im_rescale(img, scales=[1024])
+            im_batch += scaled_imgs
+            im_scales += scales
+            im_shapes += [x.shape[:2] for x in scaled_imgs]
+        im_batch = im_vstack(im_batch, self.model.pixel_mean_value, size=(1024, 1024))
+        im_shapes = np.array(im_shapes)
+        im_scales = np.array(im_scales).reshape((len(im_batch), -1))
+        im_info = np.hstack([im_shapes, im_scales]).astype("float32")
+        return im_batch, im_info
+
+    @torch.inference_mode()
+    def get_results(self, examples):
+        """Return the results."""
+        # Preprocess images and prompts.
+        imgs = [example["img"] for example in examples]
+        points = np.concatenate([example["points"] for example in examples])
+        im_batch, im_info = self.preprocess_images(imgs)
+        num_prompts = points.shape[0] if len(points.shape) > 2 else 1
+        batch_shape = im_batch.shape[0], num_prompts // im_batch.shape[0]
+        batch_points = points.reshape(batch_shape + (-1, 3))
+        batch_points[:, :, :, :2] *= im_info[:, None, None, 2:4]
+        batch_points = batch_points.reshape(points.shape)
+        # Predict tokens and masks.
+        inputs = self.model.get_inputs({"img": im_batch})
+        inputs.update(self.model.get_features(inputs))
+        outputs = self.model.get_outputs(dict(**inputs, **{"points": batch_points}))
+        # Select final mask.
+        iou_pred = outputs["iou_pred"].cpu().numpy()
+        point_score = batch_points[:, 0, 2].__eq__(2).__sub__(0.5)[:, None]
+        rank_scores = iou_pred + point_score * ([1000] + [0] * (iou_pred.shape[1] - 1))
+        mask_index = np.arange(rank_scores.shape[0]), rank_scores.argmax(1)
+        iou_scores = outputs["iou_pred"][mask_index].cpu().numpy().reshape(batch_shape)
+        # Upscale masks to the original image resolution.
+        mask_pred = outputs["mask_pred"][mask_index][:, None]
+        mask_pred = self.model.upscale_masks(mask_pred, im_batch.shape[1:-1])
+        mask_pred = mask_pred.view(batch_shape + mask_pred.shape[2:])
+        # Predict concepts.
+        concepts, scores = self.model.predict_concept(outputs["sem_embeds"][mask_index])
+        concepts, scores = [x.reshape(batch_shape) for x in (concepts, scores)]
+        # Generate captions.
+        sem_tokens = outputs["sem_tokens"][mask_index][:, None, :]
+        captions = self.model.generate_text(sem_tokens).reshape(batch_shape)
+        # Postprecess results.
+        results = []
+        for i in range(batch_shape[0]):
+            pred_h, pred_w = im_info[i, :2].astype("int")
+            masks = mask_pred[i : i + 1, :, :pred_h, :pred_w]
+            masks = self.model.upscale_masks(masks, imgs[i].shape[:2])[0]
+            results.append(
+                {
+                    "scores": np.stack([iou_scores[i], scores[i]], axis=-1),
+                    "masks": masks.gt(0).cpu().numpy().astype("uint8"),
+                    "concepts": concepts[i],
+                    "captions": captions[i],
+                }
+            )
+        return results
+
+
+class ServingCommand(object):
+    """Command to run serving."""
+
+    def __init__(self, output_queue):
+        self.output_queue = output_queue
+        self.output_dict = mp.Manager().dict()
+        self.output_index = mp.Value("i", 0)
+
+    def postprocess_outputs(self, outputs):
+        """Main the detection objects."""
+        scores, masks = outputs["scores"], outputs["masks"]
+        concepts, captions = outputs["concepts"], outputs["captions"]
+        text_template = "{} ({:.2f}, {:.2f}): {}"
+        text_contents = concepts, scores[:, 0], scores[:, 1], captions
+        texts = np.array([text_template.format(*vals) for vals in zip(*text_contents)])
+        return masks, texts
+
+    def run(self):
+        """Main loop to make the serving outputs."""
+        while True:
+            img_id, outputs = self.output_queue.get()
+            self.output_dict[img_id] = self.postprocess_outputs(outputs)
+
+
+def build_gradio_app(queues, command):
+    """Build the gradio application."""
+    import cv2
+    import gradio as gr
+    import gradio_image_prompter as gr_ext
+
+    title = "Tokenize Anything"
+    header = (
+        "<div align='center'>"
+        f"<h1>{title}</h1>"
+        "<h3>A promptable model capable of simultaneously segmenting, recognizing and captioning</h3>"
+        "</div>"
+    )
+    theme = "soft"
+    css = """#anno-img .mask {opacity: 0.5; transition: all 0.2s ease-in-out;}
+             #anno-img .mask.active {opacity: 0.7}"""
+
+    def get_examples():
+        assets_dir = os.path.join(os.path.dirname(__file__), "../assets")
+        app_images = list(filter(lambda x: x.startswith("app_image"), os.listdir(assets_dir)))
+        app_images.sort()
+        return [{"image": os.path.join(assets_dir, x)} for x in app_images]
+
+    def on_prompt_opt(index):
+        click_img = gr.Image(None, visible=index == 0)
+        draw_img = gr.ImageEditor(None, visible=index != 0)
+        anno_img = gr.AnnotatedImage(None)
+        return click_img, draw_img, anno_img
+
+    def on_reset_btn():
+        click_img, draw_img = gr.Image(None), gr.ImageEditor(None)
+        anno_img = gr.AnnotatedImage(None)
+        return click_img, draw_img, anno_img
+
+    def on_submit_btn(click_img, mask_img, prompt, multipoint):
+        if prompt == 0:
+            img = cv2.imread(click_img["image"])
+            points = np.array(click_img["points"]).reshape((-1, 2, 3))
+            if multipoint == 1:
+                points = points.reshape((-1, 3))
+                lt = points[np.where(points[:, 2] == 2)[0]][None, :, :]
+                rb = points[np.where(points[:, 2] == 3)[0]][None, :, :]
+                poly = points[np.where(points[:, 2] <= 1)[0]][None, :, :]
+                points = [lt, rb, poly] if len(lt) > 0 else [poly, np.array([[[0, 0, 4]]])]
+                points = np.concatenate(points, axis=1)
+            points = (np.array([[[0, 0, 4]]]) if len(points) == 0 else points).astype("float32")
+        elif prompt == 1:
+            img, points = mask_img["background"][:, :, (2, 1, 0)], []
+            for layer in mask_img["layers"]:
+                ys, xs = np.nonzero(layer[:, :, 0])
+                keep = np.linspace(0, ys.shape[0], 11, dtype="int64")[1:-1]
+                points.append(np.stack([xs[keep][None, :], ys[keep][None, :]], 2))
+            points = np.concatenate(points).astype("float32")
+            points = np.pad(points, [(0, 0), (0, 0), (0, 1)], constant_values=1)
+            pad_points = np.array([[[0, 0, 4]]], "float32").repeat(points.shape[0], 0)
+            points = np.concatenate([points, pad_points], axis=1)
+        inputs = {"img": img, "points": points}
+        with command.output_index.get_lock():
+            command.output_index.value += 1
+            img_id = command.output_index.value
+        queues[img_id % len(queues)].put((img_id, inputs))
+        while img_id not in command.output_dict:
+            time.sleep(0.005)
+        masks, texts = command.output_dict.pop(img_id)
+        annotations = [(x, y) for x, y in zip(masks, texts)]
+        return inputs["img"][:, :, ::-1], annotations
+
+    app = gr.Blocks(title=title, theme=theme, css=css).__enter__()
+    gr.Markdown(header)
+    container, column = gr.Row().__enter__(), gr.Column().__enter__()
+    click_img = gr_ext.ImagePrompter(type="filepath", show_label=False)
+    draw_img = gr.ImageEditor(type="numpy", show_label=False, visible=False)
+    interactions = "LeftClick (FG) | MiddleClick (BG) | PressMove (Box) | Draw (Sketch)"
+    gr.Markdown("<h3 style='text-align: center'>[🖱️ | 🖐️]: 🌟🌟 {} 🌟🌟 </h3>".format(interactions))
+    row = gr.Row().__enter__()
+    prompt_opt = gr.Radio(["Point+Box", "Sketch"], label="Prompt", type="index", value="Point+Box")
+    point_opt = gr.Radio(["Batch", "Ensemble"], label="Multipoint", type="index", value="Batch")
+    _, row = row.__exit__(), gr.Row().__enter__()
+    reset_btn, submit_btn = gr.Button("Reset"), gr.Button("Execute")
+    _, row = row.__exit__(), gr.Row().__enter__()
+    gr.Examples(get_examples(), inputs=[click_img], label="Examples (for Point+Box only)")
+    _, _, column = row.__exit__(), column.__exit__(), gr.Column().__enter__()
+    anno_img = gr.AnnotatedImage(elem_id="anno-img", show_label=False)
+    reset_btn.click(on_reset_btn, [], [click_img, draw_img, anno_img])
+    submit_btn.click(on_submit_btn, [click_img, draw_img, prompt_opt, point_opt], [anno_img])
+    prompt_opt.change(on_prompt_opt, [prompt_opt], [click_img, draw_img, anno_img])
+    column.__exit__(), container.__exit__(), app.__exit__()
+    return app
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    queues = [mp.Queue(1024) for _ in range(len(args.device) + 1)]
+    commands = [
+        test_engine.InferenceCommand(
+            queues[i],
+            queues[-1],
+            kwargs={
+                "model_type": args.model_type,
+                "weights": args.checkpoint,
+                "concept_weights": args.concept,
+                "device": args.device[i],
+                "predictor_type": Predictor,
+                "verbose": i == 0,
+            },
+        )
+        for i in range(len(args.device))
+    ]
+    commands += [ServingCommand(queues[-1])]
+    actors = [mp.Process(target=command.run, daemon=True) for command in commands]
+    for actor in actors:
+        actor.start()
+    app = build_gradio_app(queues[:-1], commands[-1])
+    app.queue()
+    app.launch()

concepts/merged_2560.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7a17403190a7a44669136d0ab278b1bb1e095bb68eff178c3e2617b2744bbb7
+size 10514948

models/tap_vit_l_03f8ec.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d63a5aba993c34bf29c0466026136e18e25d2bd4ac9e51b8fc407b76c431707d
+size 811637521

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python
+Pillow
+gradio-image-prompter
+torch>=2.0.0
+flash-attn>=2.3.3

tokenize_anything/__init__.py

@@ -0,0 +1,19 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Simultaneously Segment, Recognize, and Caption Anything with Promptable Tokenization."""
+
+from tokenize_anything.build_model import model_registry
+from tokenize_anything.version import __version__

tokenize_anything/build_model.py

@@ -0,0 +1,114 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Build model."""
+
+from functools import partial
+import pickle
+
+import numpy as np
+import torch
+
+from tokenize_anything.modeling import ConceptProjector
+from tokenize_anything.modeling import ImageDecoder
+from tokenize_anything.modeling import ImageEncoderViT
+from tokenize_anything.modeling import ImageTokenizer
+from tokenize_anything.modeling import PromptEncoder
+from tokenize_anything.modeling import TextDecoder
+from tokenize_anything.modeling import TextTokenizer
+
+
+def get_device(device_index):
+    """Create an available device object."""
+    if torch.cuda.is_available():
+        return torch.device("cuda", device_index)
+    return torch.device("cpu")
+
+
+def load_weights(module, weights_file, strict=True):
+    """Load a weights file."""
+    if not weights_file:
+        return module._IncompatibleKeys([], [])
+    if weights_file.endswith(".pkl"):
+        with open(weights_file, "rb") as f:
+            state_dict = pickle.load(f)
+            for k, v in state_dict.items():
+                state_dict[k] = torch.from_numpy(v) if isinstance(v, np.ndarray) else v
+    else:
+        state_dict = torch.load(weights_file)
+    return module.load_state_dict(state_dict, strict=strict)
+
+
+def vit_encoder(depth, embed_dim, num_heads, out_dim, image_size):
+    """Build an image encoder with ViT."""
+    return ImageEncoderViT(
+        depth=depth,
+        embed_dim=embed_dim,
+        num_heads=num_heads,
+        mlp_ratio=4,
+        patch_size=16,
+        window_size=16,
+        image_size=image_size,
+        out_dim=out_dim,
+    )
+
+
+def image_tokenizer(image_encoder, checkpoint=None, device=0, dtype="float16", **kwargs):
+    """Build an image tokenizer."""
+    image_size = kwargs.get("image_size", 1024)
+    prompt_embed_dim = kwargs.get("prompt_embed_dim", 256)
+    sem_embed_dim = kwargs.get("sem_embed_dim", 1024)
+    text_embed_dim = kwargs.get("text_embed_dim", 512)
+    text_decoder_depth = kwargs.get("text_decoder_depth", 12)
+    text_seq_len = kwargs.get("text_seq_len", 40)
+    text_tokenizer = TextTokenizer()
+    model = ImageTokenizer(
+        image_encoder=image_encoder(out_dim=prompt_embed_dim, image_size=image_size),
+        prompt_encoder=PromptEncoder(embed_dim=prompt_embed_dim, image_size=image_size),
+        image_decoder=ImageDecoder(
+            depth=2,
+            embed_dim=prompt_embed_dim,
+            num_heads=prompt_embed_dim // 32,
+            num_mask_tokens=4,
+            sem_embed_dim=sem_embed_dim,
+        ),
+        text_tokenizer=text_tokenizer,
+        concept_projector=ConceptProjector(),
+        text_decoder=TextDecoder(
+            depth=text_decoder_depth,
+            embed_dim=text_embed_dim,
+            num_heads=text_embed_dim // 64,
+            mlp_ratio=4,
+            prompt_embed_dim=prompt_embed_dim,
+            max_seq_len=text_seq_len,
+            vocab_size=text_tokenizer.n_words,
+        ),
+    )
+    load_weights(model, checkpoint)
+    model = model.to(device=get_device(device))
+    model = model.eval() if not kwargs.get("training", False) else model
+    model = model.half() if dtype == "float16" else model
+    model = model.bfloat16() if dtype == "bfloat16" else model
+    model = model.float() if dtype == "float32" else model
+    return model
+
+
+vit_b_encoder = partial(vit_encoder, depth=12, embed_dim=768, num_heads=12)
+vit_l_encoder = partial(vit_encoder, depth=24, embed_dim=1024, num_heads=16)
+
+model_registry = {
+    "tap_vit_b": partial(image_tokenizer, image_encoder=vit_b_encoder),
+    "tap_vit_l": partial(image_tokenizer, image_encoder=vit_l_encoder),
+}

tokenize_anything/modeling/__init__.py

@@ -0,0 +1,24 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Modeling components."""
+
+from tokenize_anything.modeling.concept_projector import ConceptProjector
+from tokenize_anything.modeling.image_decoder import ImageDecoder
+from tokenize_anything.modeling.image_encoder import ImageEncoderViT
+from tokenize_anything.modeling.image_tokenizer import ImageTokenizer
+from tokenize_anything.modeling.prompt_encoder import PromptEncoder
+from tokenize_anything.modeling.text_decoder import TextDecoder
+from tokenize_anything.modeling.text_tokenizer import TextTokenizer

tokenize_anything/modeling/concept_projector.py

@@ -0,0 +1,74 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Concet projector."""
+
+import pickle
+
+import numpy as np
+import torch
+from torch import nn
+
+
+class ConceptProjector(nn.Module):
+    """Encode and decode concept using CLIP."""
+
+    def __init__(self, src_weights=None, tgt_weights=None):
+        super(ConceptProjector, self).__init__()
+        self.reset_weights(src_weights, tgt_weights)
+
+    def reset_weights(self, src_weights=None, tgt_weights=None):
+        """Reset the normalized projection weights."""
+        if src_weights is not None:
+            with open(src_weights, "rb") as f:
+                self.src_weights, self.concepts = pickle.load(f)
+                self.src_weights = torch.from_numpy(self.src_weights)
+                self.concepts = np.array(self.concepts)
+        if tgt_weights is not None:
+            with open(tgt_weights, "rb") as f:
+                self.tgt_weights, self.concepts = pickle.load(f)
+                self.tgt_weights = torch.from_numpy(self.tgt_weights)
+                self.concepts = np.array(self.concepts)
+
+    @staticmethod
+    def maybe_convert(embeds, proj):
+        """Convert inputs for safe projection."""
+        if embeds.dtype != torch.float32:
+            embeds = embeds.float()
+        if embeds.device != proj.device:
+            proj = proj.to(device=embeds.device)
+        return embeds, proj
+
+    def encode_src(self, src_embeds):
+        """Encode source visual embedding via concept projection."""
+        src_embeds, self.src_weights = self.maybe_convert(src_embeds, self.src_weights)
+        logits = nn.functional.normalize(src_embeds, dim=-1) @ self.src_weights
+        return nn.functional.log_softmax(logits, dim=-1)
+
+    def encode_tgt(self, tgt_embeds):
+        """Encode target visual embedding via concept projection."""
+        tgt_embeds, self.tgt_weights = self.maybe_convert(tgt_embeds, self.tgt_weights)
+        logits = nn.functional.normalize(tgt_embeds, dim=-1) @ self.tgt_weights
+        return nn.functional.log_softmax(logits, dim=-1)
+
+    def decode(self, src_embeds, k=1, return_index=False, return_prob=False):
+        """Return the top-k concepts of source visual embedding."""
+        src_embeds, self.src_weights = self.maybe_convert(src_embeds, self.src_weights)
+        logits = nn.functional.normalize(src_embeds, dim=-1) @ self.src_weights
+        probs = nn.functional.softmax(logits, dim=-1)
+        if return_prob:
+            return probs.cpu().numpy()
+        score, index = [x.cpu().numpy() for x in probs.topk(k, -1)]
+        return (index if return_index else self.concepts[index]), score

tokenize_anything/modeling/image_decoder.py

@@ -0,0 +1,224 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Image decoder."""
+
+try:
+    from flash_attn import flash_attn_func
+except ImportError:
+    flash_attn_func = None
+
+import torch
+from torch import nn
+
+
+class TransposedLayerNorm(nn.LayerNorm):
+    """LayerNorm with pre-transposed spatial axes."""
+
+    def forward(self, input):
+        return super().forward(input.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+
+
+class MLP(nn.Module):
+    """Two layers MLP."""
+
+    def __init__(self, dim, mlp_dim, activation_type="ReLU"):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(dim, mlp_dim)
+        self.fc2 = nn.Linear(mlp_dim, dim)
+        self.activation = getattr(nn, activation_type)()
+        self.activation.inplace = True
+
+    def forward(self, x):
+        return self.fc2(self.activation(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    """Multi-head attention."""
+
+    def __init__(self, dim=256, num_heads=8, attn_ratio=1):
+        super(Attention, self).__init__()
+        qkv_dim = int(dim * attn_ratio)
+        self.num_heads = num_heads
+        self.head_dim = qkv_dim // num_heads
+        self.q_proj = nn.Linear(dim, qkv_dim)
+        self.k_proj = nn.Linear(dim, qkv_dim)
+        self.v_proj = nn.Linear(dim, qkv_dim)
+        self.proj = nn.Linear(qkv_dim, dim)
+        self.scale = self.head_dim**-0.5
+
+    def forward(self, q, k, v):
+        q = self.q_proj(q).view((-1, q.size(1), self.num_heads, self.head_dim))
+        k = self.k_proj(k).view((-1, k.size(1), self.num_heads, self.head_dim))
+        v = self.v_proj(v).view((-1, v.size(1), self.num_heads, self.head_dim))
+        o = flash_attn_func(q, k, v, softmax_scale=self.scale)
+        return self.proj(o.flatten(2))
+
+
+class Block(nn.Module):
+    """Transformer block."""
+
+    def __init__(
+        self,
+        dim=256,
+        num_heads=8,
+        attn_ratio=0.5,
+        mlp_dim=2048,
+        dropout=0.1,
+        activation_type="ReLU",
+        skip_first_query_pos=False,
+    ):
+        super(Block, self).__init__()
+        self.self_attn = Attention(dim, num_heads)
+        self.norm1 = nn.LayerNorm(dim)
+        self.cross_attn_token_to_image = Attention(dim, num_heads, attn_ratio)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = MLP(dim, mlp_dim, activation_type)
+        self.norm3 = nn.LayerNorm(dim)
+        self.cross_attn_image_to_token = Attention(dim, num_heads, attn_ratio)
+        self.norm4 = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout, inplace=True)
+        self.skip_first_query_pos = skip_first_query_pos
+
+    def forward(self, query, key, query_pos, key_pos):
+        if self.skip_first_query_pos:
+            query = self.norm1(self.self_attn(query, query, query))
+        else:
+            q = query + query_pos
+            query = self.norm1(self.dropout(self.self_attn(q, q, query)).add_(query))
+        q, k = query + query_pos, key + key_pos
+        query = self.norm2(self.dropout(self.cross_attn_token_to_image(q, k, key)).add_(query))
+        query = self.norm3(self.dropout(self.mlp(query)).add_(query))
+        q = query + query_pos
+        key = self.norm4(self.cross_attn_image_to_token(k, q, query).add_(key))
+        return query, key
+
+
+class Transformer(nn.Module):
+    """Two-way transformer decoder."""
+
+    def __init__(
+        self,
+        embed_dim=256,
+        num_heads=8,
+        attn_ratio=0.5,
+        mlp_dim=2048,
+        dropout=0.1,
+        activation_type="ReLU",
+        depth=2,
+    ):
+        super(Transformer, self).__init__()
+        self.blocks = nn.ModuleList(
+            Block(
+                embed_dim,
+                num_heads,
+                attn_ratio=attn_ratio,
+                mlp_dim=mlp_dim,
+                dropout=dropout,
+                activation_type=activation_type,
+                skip_first_query_pos=i == 0,
+            )
+            for i in range(depth)
+        )
+        self.final_attn_token_to_image = Attention(embed_dim, num_heads, attn_ratio)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.dropout = nn.Dropout(dropout, inplace=True)
+
+    def forward(self, query, key, query_pos, key_pos):
+        for blk in self.blocks:
+            query, key = blk(query, key, query_pos, key_pos)
+        q, k = query + query_pos, key + key_pos
+        query = self.dropout(self.final_attn_token_to_image(q, k, key)).add_(query)
+        query = self.norm(query)
+        return query, key
+
+
+class Predictor(nn.Module):
+    """MLP predictor."""
+
+    def __init__(self, in_dim, out_dim, mlp_dim=None, depth=3):
+        super(Predictor, self).__init__()
+        mlp_dims = [mlp_dim or in_dim] * (depth - 1)
+        in_dims, out_dims = [in_dim] + mlp_dims, mlp_dims + [out_dim]
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip(in_dims, out_dims))
+
+    def forward(self, x):
+        for fc in self.layers[:-1]:
+            x = nn.functional.relu(fc(x), inplace=True)
+        return self.layers[-1](x)
+
+
+class ImageDecoder(nn.Module):
+    """Module to decode region tokens and masks."""
+
+    def __init__(self, depth, embed_dim, num_heads, num_mask_tokens=4, sem_embed_dim=1024):
+        super(ImageDecoder, self).__init__()
+        self.embed_dim = embed_dim
+        self.num_mask_tokens = num_mask_tokens
+        self.transformer = Transformer(embed_dim, num_heads=num_heads, depth=depth)
+        self.iou_token = nn.Embedding(1, embed_dim)
+        self.sem_tokens = nn.Embedding(self.num_mask_tokens, embed_dim)
+        self.mask_tokens = nn.Embedding(self.num_mask_tokens, embed_dim)
+        self.output_conv = nn.Sequential(
+            nn.ConvTranspose2d(embed_dim, embed_dim // 4, 2, 2),
+            TransposedLayerNorm(embed_dim // 4),
+            nn.GELU(),
+            nn.ConvTranspose2d(embed_dim // 4, embed_dim // 8, 2, 2),
+            nn.GELU(),
+        )
+        self.mask_pred = nn.ModuleList(
+            Predictor(embed_dim, embed_dim // 8) for _ in range(num_mask_tokens)
+        )
+        self.iou_pred = Predictor(embed_dim, self.num_mask_tokens)
+        self.sem_pred = Predictor(embed_dim, sem_embed_dim, 1024)
+
+    def get_outputs(self, inputs):
+        img_embeds = inputs["img_embeds"]
+        sparse_embeds = inputs["sparse_embeds"]
+        ims_per_batch = img_embeds.size(0)
+        prompts_per_batch = sparse_embeds.size(0)
+        img_embed_size = img_embeds.shape[2:-1]
+        # Prepare query.
+        tokens = [self.sem_tokens.weight, self.iou_token.weight, self.mask_tokens.weight]
+        query = torch.cat(tokens).unsqueeze_(0).expand(prompts_per_batch, -1, -1)
+        query = torch.cat((query, sparse_embeds), dim=1)
+        num_tokens = query.shape[1] - sparse_embeds.shape[1]
+        # Prepare key.
+        key = img_embeds.expand(-1, prompts_per_batch // ims_per_batch, -1, -1, -1)
+        key = key.flatten(0, 1).flatten(1, 2)
+        # Decode.
+        query, key = self.transformer(query, key, query, inputs["img_pos"])
+        # Upscale key.
+        key = key.transpose(1, 2).view((-1, self.embed_dim) + img_embed_size)
+        output_masks = self.output_conv(key).flatten(2)
+        # Unpack query.
+        tokens = query[:, :num_tokens].unbind(dim=1)
+        iou_tokens = tokens[num_tokens - self.num_mask_tokens - 1]
+        mask_tokens = tokens[num_tokens - self.num_mask_tokens :]
+        sem_tokens = tokens[: self.num_mask_tokens]
+        # Predict.
+        mask_pred = [f(x) for f, x in zip(self.mask_pred, mask_tokens)]
+        mask_pred = torch.stack(mask_pred, dim=1) @ output_masks
+        mask_pred_size = list(4 * embed_size for embed_size in img_embed_size)
+        mask_pred = mask_pred.view([-1, self.num_mask_tokens] + mask_pred_size)
+        outputs = {"iou_pred": self.iou_pred(iou_tokens), "mask_pred": mask_pred}
+        outputs["sem_tokens"] = torch.stack(sem_tokens, dim=1)
+        outputs["sem_embeds"] = self.sem_pred(outputs["sem_tokens"])
+        return outputs
+
+    def forward(self, inputs):
+        outputs = self.get_outputs(inputs)
+        outputs["iou_pred"] = outputs["iou_pred"].float()
+        return outputs

tokenize_anything/modeling/image_encoder.py

@@ -0,0 +1,254 @@
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+##############################################################################
+"""Image encoder."""
+
+import torch
+from torch import nn
+
+
+def space_to_depth(input, block_size):
+    """Rearrange blocks of spatial data into depth."""
+    if input.dim() == 3:
+        hXw, c = input.size()[1:]
+        h = w = int(hXw**0.5)
+    else:
+        h, w, c = input.size()[1:]
+    h1, w1 = h // block_size, w // block_size
+    c1 = (block_size**2) * c
+    input = input.reshape((-1, h1, block_size, w1, block_size, c))
+    return input.permute(0, 1, 3, 2, 4, 5).reshape((-1, h1, w1, c1))
+
+
+def depth_to_space(input, block_size):
+    """Rearrange blocks of depth data into spatial."""
+    h1, w1, c1 = input.size()[1:]
+    h, w = h1 * block_size, w1 * block_size
+    c = c1 // (block_size**2)
+    input = input.reshape((-1, h1, w1, block_size, block_size, c))
+    return input.permute(0, 1, 3, 2, 4, 5).reshape((-1, h, w, c))
+
+
+class MLP(nn.Module):
+    """Two layers MLP."""
+
+    def __init__(self, dim, mlp_ratio=4):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(dim, int(dim * mlp_ratio))
+        self.fc2 = nn.Linear(int(dim * mlp_ratio), dim)
+        self.activation = nn.GELU()
+
+    def forward(self, x):
+        return self.fc2(self.activation(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    """Multihead attention."""
+
+    def __init__(self, dim, num_heads, qkv_bias=True):
+        super(Attention, self).__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+        self.rel_pos_embed = nn.Identity()
+
+    def forward(self, x):
+        qkv_shape = (-1, x.size(1), 3, self.num_heads, self.head_dim)
+        qkv = self.qkv(x).reshape(qkv_shape).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(dim=0)
+        attn = q @ k.transpose(-2, -1).mul(self.scale)
+        attn = self.rel_pos_embed(attn)
+        o = nn.functional.softmax(attn, dim=-1) @ v
+        return self.proj(o.transpose(1, 2).flatten(2))
+
+
+class Block(nn.Module):
+    """Transformer block."""
+
+    def __init__(self, dim, num_heads, mlp_ratio=4, qkv_bias=True):
+        super(Block, self).__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = Attention(dim, num_heads, qkv_bias=qkv_bias)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = MLP(dim, mlp_ratio=mlp_ratio)
+
+    def forward(self, x):
+        x = self.attn(self.norm1(x)).add_(x)
+        return self.mlp(self.norm2(x)).add_(x)
+
+
+class Bottleneck(nn.Module):
+    """The bottleneck block."""
+
+    def __init__(self, dim, expansion=2, width=None):
+        super(Bottleneck, self).__init__()
+        width = width or dim // expansion
+        self.conv1 = nn.Conv2d(dim, width, 1, bias=False)
+        self.norm1 = nn.SyncBatchNorm(width)
+        self.conv2 = nn.Conv2d(width, width, 3, padding=1, bias=False)
+        self.norm2 = nn.SyncBatchNorm(width)
+        self.conv3 = nn.Conv2d(width, dim, 1, bias=False)
+        self.norm3 = nn.SyncBatchNorm(dim)
+        self.activation = nn.GELU()
+
+    def forward(self, x):
+        shortcut = x
+        x = self.activation(self.norm1(self.conv1(x)))
+        x = self.activation(self.norm2(self.conv2(x)))
+        return self.norm3(self.conv3(x)).add_(shortcut)
+
+
+class PatchEmbed(nn.Module):
+    """Patch embedding layer."""
+
+    def __init__(self, dim=768, patch_size=16, bias=True):
+        super(PatchEmbed, self).__init__()
+        self.proj = nn.Conv2d(3, dim, patch_size, patch_size, bias=bias)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class PosEmbed(nn.Module):
+    """Position embedding layer."""
+
+    def __init__(self, dim, num_patches):
+        super(PosEmbed, self).__init__()
+        self.dim = dim
+        self.num_patches = num_patches
+        self.weight = nn.Parameter(torch.zeros(num_patches, dim))
+        nn.init.normal_(self.weight, std=0.02)
+
+    def forward(self, x):
+        return x.add_(self.weight)
+
+
+class RelPosEmbed(nn.Module):
+    """Relative position embedding layer."""
+
+    def __init__(self, num_heads, size):
+        super(RelPosEmbed, self).__init__()
+        self.register_buffer("index", self.get_index(size))
+        self.weight = nn.Parameter(torch.zeros(num_heads, (2 * size - 1) ** 2))
+
+    @staticmethod
+    def get_index(size):
+        """Return the relative index."""
+        grid = torch.arange(size)
+        grid = torch.stack(torch.meshgrid(grid, grid, indexing="ij")).reshape((2, -1))
+        coords = grid[:, :, None] - grid[:, None, :] + (size - 1)
+        coords[0] *= 2 * size - 1
+        return coords.sum(0)
+
+    def get_bias(self):
+        return self.weight[:, self.index]
+
+    def forward(self, x):
+        return x.add_(self.get_bias())
+
+
+class SimpleFeaturePyramid(nn.Module):
+    """Module to create pyramid features."""
+
+    def __init__(self, embed_dim, out_dim, patch_size=16, min_lvl=4, max_lvl=4):
+        super(SimpleFeaturePyramid, self).__init__()
+        self.min_lvl, self.max_lvl = min_lvl, max_lvl
+        self.input_conv = nn.ModuleList()
+        self.lateral_conv = nn.ModuleList()
+        self.output_conv = nn.ModuleList()
+        patch_lvl = dict((2**i, i) for i in range(6))[patch_size]
+        for lvl in [min(i + 2, self.max_lvl) for i in range(4)]:
+            if lvl == patch_lvl or lvl < self.min_lvl:
+                self.input_conv += [nn.Identity()]
+            elif lvl < patch_lvl:
+                stride, layers = 2 ** (patch_lvl - lvl), []
+                while stride > 1:
+                    layers += [nn.ConvTranspose2d(embed_dim, embed_dim, 2, 2)]
+                    layers += [nn.SyncBatchNorm(embed_dim), nn.GELU()] if stride > 2 else []
+                    stride /= 2
+                self.input_conv.append(nn.Sequential(*layers))
+            elif lvl > patch_lvl:
+                stride = 2 ** (lvl - patch_lvl)
+                self.input_conv += [nn.MaxPool2d(stride, stride)]
+        for _ in range(min_lvl, max_lvl + 1):
+            self.lateral_conv.append(
+                nn.Sequential(
+                    nn.Conv2d(embed_dim, out_dim, kernel_size=1, bias=False),
+                    nn.SyncBatchNorm(out_dim),
+                )
+            )
+            self.output_conv.append(
+                nn.Sequential(
+                    nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1, bias=False),
+                    nn.SyncBatchNorm(out_dim),
+                )
+            )
+
+    def forward(self, inputs):
+        inputs = inputs + [inputs[-1]] * (4 - len(inputs))
+        inputs = [conv(x) for conv, x in zip(self.input_conv, inputs)]
+        features = inputs[self.min_lvl - 1 : self.max_lvl]
+        laterals = [conv(x) for conv, x in zip(self.lateral_conv, features)]
+        return [conv(x) for conv, x in zip(self.output_conv, laterals)]
+
+
+class ImageEncoderViT(nn.Module):
+    """ViT image encoder."""
+
+    def __init__(
+        self,
+        depth,
+        embed_dim,
+        num_heads,
+        mlp_ratio=4,
+        patch_size=16,
+        window_size=16,
+        image_size=1024,
+        out_dim=256,
+    ):
+        super(ImageEncoderViT, self).__init__()
+        self.embed_dim = embed_dim
+        self.image_size = image_size
+        self.window_size = window_size or image_size // patch_size
+        self.patch_embed = PatchEmbed(embed_dim, patch_size)
+        self.pos_embed = PosEmbed(embed_dim, (image_size // patch_size) ** 2)
+        self.blocks = nn.ModuleList(Block(embed_dim, num_heads, mlp_ratio) for _ in range(depth))
+        for blk in self.blocks:
+            blk.attn.rel_pos_embed = RelPosEmbed(num_heads, self.window_size)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.cross_conv = nn.ModuleList(Bottleneck(embed_dim) for _ in range(4))
+        self.neck = SimpleFeaturePyramid(embed_dim, out_dim, patch_size)
+        self.cross_indices = list(range(depth // 4 - 1, depth, depth // 4))
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        x = self.pos_embed(x)
+        x = space_to_depth(x, self.window_size)
+        wmsa_shape = (-1,) + x.shape[1:]
+        msa_shape = (-1, self.window_size**2, self.embed_dim)
+        x = x.reshape(msa_shape)
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in self.cross_indices or i == len(self.blocks) - 1:
+                x = self.norm(x) if i == len(self.blocks) - 1 else x
+                x = depth_to_space(x.reshape(wmsa_shape), self.window_size)
+                x = x.permute(0, 3, 1, 2)
+            if i in self.cross_indices:
+                x = self.cross_conv[self.cross_indices.index(i)](x)
+            if i in self.cross_indices and i < len(self.blocks) - 1:
+                x = x.permute(0, 2, 3, 1)
+                x = space_to_depth(x, self.window_size).reshape(msa_shape)
+        return self.neck([x])

tokenize_anything/modeling/image_tokenizer.py

@@ -0,0 +1,201 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Image tokenizer."""
+
+import numpy as np
+import torch
+from torch import nn
+
+
+class ImageTokenizer(nn.Module):
+    """Tokenize image regions with visual prompts."""
+
+    def __init__(
+        self,
+        image_encoder,
+        prompt_encoder,
+        image_decoder,
+        concept_projector=None,
+        text_tokenizer=None,
+        text_decoder=None,
+        pixel_mean=(103.53, 116.28, 123.675),
+        pixel_std=(57.375, 57.12, 58.395),
+    ):
+        super(ImageTokenizer, self).__init__()
+        self.image_encoder = image_encoder
+        self.prompt_encoder = prompt_encoder
+        self.image_decoder = image_decoder
+        self.concept_projector = concept_projector
+        self.text_tokenizer = text_tokenizer
+        self.text_decoder = text_decoder
+        self.pixel_mean_value = pixel_mean  # BGR order.
+        self.register_buffer("pixel_mean", torch.Tensor(pixel_mean))
+        self.register_buffer("pixel_rsig", torch.Tensor(pixel_std).reciprocal_())
+
+    def get_inputs(self, inputs):
+        """Return the model inputs.
+
+        Parameters
+        ----------
+        inputs : dict
+            The initial inputs.
+
+        Returns
+        -------
+        dict
+            The model inputs.
+
+        """
+        if not isinstance(inputs["img"], torch.Tensor):
+            inputs["img"] = torch.from_numpy(inputs["img"])
+        if inputs["img"].device != self.pixel_mean.device:
+            inputs["img"] = inputs["img"].to(device=self.pixel_mean.device)
+        inputs["img"] = inputs["img"].to(dtype=self.pixel_mean.dtype)
+        inputs["img"] = inputs["img"].sub(self.pixel_mean).mul_(self.pixel_rsig)
+        inputs["img"] = inputs["img"].permute(0, 3, 1, 2)
+        return inputs
+
+    def get_features(self, inputs):
+        """Return the image features.
+
+        Parameters
+        ----------
+        inputs : dict
+            The inputs.
+
+        Returns
+        -------
+        dict
+            The image features.
+
+        """
+        features = self.image_encoder(inputs["img"])
+        img_embeds = features[0].permute(0, 2, 3, 1).unsqueeze_(1)
+        return {"features": features, "img_embeds": img_embeds}
+
+    def get_outputs(self, inputs):
+        """Return the model outputs.
+
+        Parameters
+        ----------
+        inputs : dict
+            The model inputs.
+
+        Returns
+        -------
+        dict
+            The model outputs.
+
+        """
+        inputs.update(self.prompt_encoder(inputs))
+        return self.image_decoder(inputs)
+
+    def forward(self, inputs):
+        """Define the computation performed at every call.
+
+        Parameters
+        ----------
+        inputs : dict
+            The initial inputs.
+
+        Returns
+        -------
+        dict
+            The model outputs.
+
+        """
+        inputs = self.get_inputs(inputs)
+        inputs.update(self.get_features(inputs))
+        return self.get_outputs(inputs)
+
+    def upscale_masks(self, masks, size):
+        """Upscale masks using bilinear interpolation.
+
+        Parameters
+        ----------
+        masks : torch.Tensor
+            The input masks.
+        size : Union[int, Tuple[int]]
+            The output size.
+
+        Returns
+        -------
+        torch.Tensor
+            The output masks.
+
+        """
+        return nn.functional.interpolate(masks, size, mode="bilinear", align_corners=False)
+
+    @torch.inference_mode()
+    def predict_concept(self, visual_embeds, k=1):
+        """Predict top-k concepts based on visual embeddings.
+
+        Parameters
+        ----------
+        visual_embeds: torch.Tensor
+            The embeddings to predict visual content.
+        k : int, optional, default=1
+            The k value.
+
+        Returns
+        -------
+        Tuple[numpy.ndarray, numpy.ndarray]
+            The concept scores and indices.
+
+        """
+        return self.concept_projector.decode(visual_embeds, k)
+
+    @torch.inference_mode()
+    def generate_text(self, visual_tokens, max_gen_len=None, temperature=0):
+        """Generate text sequences based on visual tokens.
+
+        Parameters
+        ----------
+        visual_tokens: torch.Tensor
+            The tokens to prompt visual context.
+        max_gen_len : int, optional
+            The maximum length of the generated text sequences.
+        temperature : float, optional
+            The temperature for controlling randomness in sampling.
+
+        Returns
+        -------
+        np.ndarray
+            An array of generated texts.
+
+        """
+        max_gen_len = max_gen_len or self.text_decoder.max_seq_len
+        prompts = self.text_decoder.get_prompts(visual_tokens)
+        out_shape = (prompts.size(0), self.text_decoder.max_text_len)
+        tokens = np.full(out_shape, self.text_tokenizer.pad_id, "int64")
+        tokens[:, 0], prev_pos = self.text_tokenizer.bos_id, 0
+        eos_reached = np.array([False] * tokens.shape[0])
+        for cur_pos in range(1, max_gen_len):
+            decode_seq_len = cur_pos - prev_pos
+            x = torch.from_numpy(tokens[:, prev_pos:cur_pos]).to(device=prompts.device)
+            logits = self.text_decoder.transformer(prompts, x, prev_pos)
+            next_logits = logits[: x.size(0), decode_seq_len - 1]
+            if temperature > 0:
+                p = nn.functional.softmax(next_logits / temperature, dim=-1)
+                next_token = torch.multinomial(p, 1).cpu().numpy().flatten()
+            else:
+                next_token = next_logits.argmax(-1).cpu().numpy()
+            tokens[:, cur_pos] = next_token
+            eos_reached |= next_token == self.text_tokenizer.eos_id
+            prev_pos, logits, next_logits = cur_pos, None, None
+            if eos_reached.all():
+                break
+        return np.array(self.text_tokenizer.detokenize(tokens))

tokenize_anything/modeling/prompt_encoder.py

@@ -0,0 +1,100 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Prompt encoder."""
+
+import torch
+from torch import nn
+
+
+class PromptEncoder(nn.Module):
+    """Module to encode geometric prompts."""
+
+    def __init__(self, embed_dim, image_size):
+        super(PromptEncoder, self).__init__()
+        self.img_size = [image_size] * 2
+        self.point_embed = nn.Embedding(5, embed_dim)  # [bg, fg, lt, rb, pad]
+        self.corner_labels = torch.tensor([[2, 3]], dtype=torch.int64)
+        self.register_buffer("coord_matrix", torch.randn((2, embed_dim // 2)))
+        self.img_pos = None
+
+    def to_tensor(self, input):
+        """Convert input to tensor."""
+        if input is None:
+            return input
+        if not isinstance(input, torch.Tensor):
+            input = torch.from_numpy(input)
+        if input.device != self.coord_matrix.device:
+            input = input.to(device=self.coord_matrix.device)
+        return input
+
+    def to_points(self, points=None, boxes=None):
+        """Convert points or boxes to point prompts."""
+        if points is not None:
+            if isinstance(points, (tuple, list)):
+                coords, labels = points
+            else:
+                coords, labels = points[:, :, :2], points[:, :, 2]
+            coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1])
+            coords = self.to_tensor(coords.clip(0, 1).astype("float32"))
+            labels = self.to_tensor(labels.astype("int64"))
+            return coords, labels
+        if boxes is not None:
+            coords = boxes.reshape((-1, 2, 2))
+            coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1])
+            coords = self.to_tensor(coords.clip(0, 1).astype("float32"))
+            labels = self.to_tensor(self.corner_labels)
+            return coords, labels
+        return None
+
+    def encode_coords(self, coords):
+        """Return the embedding for given coords."""
+        pi4, pi2 = 4 * 3.1415926, 2 * 3.1415926
+        if self.coord_matrix.dtype != torch.float32:
+            self.coord_matrix = self.coord_matrix.float()
+        rad = coords.mul(pi4).sub_(pi2) @ self.coord_matrix
+        dtype = self.point_embed.weight.dtype
+        return torch.cat([rad.sin(), rad.cos()], dim=-1).to(dtype=dtype)
+
+    def encode_points(self, coords, labels):
+        """Return the embedding for given points."""
+        embed = self.encode_coords(coords)
+        embed.mul_(labels.ne(4).unsqueeze_(-1).float().to(dtype=embed.dtype))
+        return embed.add_(self.point_embed(labels))
+
+    def encode_grid(self, grid_size):
+        """Return the embedding for a grid of specified size."""
+        grid = torch.ones(*grid_size, dtype=torch.float32)
+        y = grid.cumsum(dim=0).sub_(0.5).div_(grid_size[0])
+        x = grid.cumsum(dim=1).sub_(0.5).div_(grid_size[1])
+        coords = self.to_tensor(torch.stack([x, y], dim=-1))
+        return self.encode_coords(coords)
+
+    def forward(self, inputs):
+        sparse_embeds = []
+        if inputs.get("boxes", None) is not None:
+            coords, labels = self.to_points(boxes=inputs["boxes"])
+            sparse_embeds.append(self.encode_points(coords, labels))
+        if inputs.get("points", None) is not None:
+            coords, labels = self.to_points(points=inputs["points"])
+            sparse_embeds.append(self.encode_points(coords, labels))
+        if len(sparse_embeds) > 1:
+            sparse_embeds = [torch.cat(sparse_embeds, dim=1)]
+        elif len(sparse_embeds) == 0:
+            raise ValueError("Excepted ``points`` or ``boxes`` prompts.")
+        img_embed_size = torch.Size(inputs["img_embeds"].shape[2:-1])
+        if self.img_pos is None or self.img_pos.shape[0] != img_embed_size.numel():
+            self.img_pos = self.encode_grid(img_embed_size).flatten(0, 1)
+        return {"sparse_embeds": sparse_embeds[0], "img_pos": self.img_pos}

tokenize_anything/modeling/text_decoder.py

@@ -0,0 +1,206 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Text decoder."""
+
+try:
+    from flash_attn import flash_attn_func
+    from flash_attn import flash_attn_with_kvcache
+    from flash_attn.layers.rotary import apply_rotary_emb
+except ImportError:
+    flash_attn_func = None
+    flash_attn_with_kvcache = None
+    apply_rotary_emb = None
+
+import torch
+from torch import nn
+
+
+class TransformerCache(nn.Module):
+    """Transformer cache module."""
+
+    def __init__(self, device=None, dtype=None):
+        super(TransformerCache, self).__init__()
+        self.device = device
+        self.dtype = dtype
+        self.start_pos = 0
+        self.cache_dict = {}
+
+    def init_seq(self, max_batch_size):
+        seq_lens = torch.zeros(max_batch_size, dtype=torch.int32, device=self.device)
+        self.cache_dict["seq_lens"] = seq_lens
+
+    def init_rotary(self, seq_len, dim, theta=10000.0):
+        grid = torch.arange(seq_len, dtype=torch.float32).unsqueeze_(-1)
+        freq = torch.pow(theta, torch.arange(0, dim, 2)[: dim // 2].float().div_(dim))
+        broadcast_freq = grid.mul(freq.reciprocal_().unsqueeze_(0))
+        cache_cos = broadcast_freq.cos().view((-1, dim // 2))
+        cache_sin = broadcast_freq.sin().view((-1, dim // 2))
+        self.cache_dict["cos"] = cache_cos.to(self.device, self.dtype)
+        self.cache_dict["sin"] = cache_sin.to(self.device, self.dtype)
+
+    def init_kv(self, mixer, kv_size):
+        cache_k = torch.zeros(*kv_size, dtype=self.dtype, device=self.device)
+        cache_v = torch.zeros(*kv_size, dtype=self.dtype, device=self.device)
+        self.cache_dict[f"{id(mixer)}_k"] = cache_k
+        self.cache_dict[f"{id(mixer)}_v"] = cache_v
+
+    def set_seq(self, start_pos=0, end_pos=None):
+        self.start_pos = start_pos
+        if "seq_lens" in self.cache_dict:
+            self.cache_dict["seq_lens"].fill_(start_pos)
+        if "cos" in self.cache_dict and end_pos is not None:
+            self.cache_dict["seq_cos"] = self.cache_dict["cos"][self.start_pos : end_pos]
+            self.cache_dict["seq_sin"] = self.cache_dict["sin"][self.start_pos : end_pos]
+
+    def forward_rotary(self, q, k, inplace=False):
+        cos = self.cache_dict.get("seq_cos", self.cache_dict.get("cos", None))
+        sin = self.cache_dict.get("seq_sin", self.cache_dict.get("sin", None))
+        if cos is None or sin is None:
+            return q, k
+        q = apply_rotary_emb(q, cos, sin, interleaved=True, inplace=inplace)
+        k = apply_rotary_emb(k, cos, sin, interleaved=True, inplace=inplace)
+        return q, k
+
+    def forward_flash(self, mixer, q, k, v):
+        cache_k = self.cache_dict.get(f"{id(mixer)}_k", None)
+        cache_v = self.cache_dict.get(f"{id(mixer)}_v", None)
+        flash_args = {"softmax_scale": mixer.scale, "causal": True}
+        if cache_k is None or cache_v is None:
+            return flash_attn_func(q, k, v, **flash_args)
+        flash_args["cache_seqlens"] = self.cache_dict["seq_lens"][: q.shape[0]]
+        return flash_attn_with_kvcache(q, cache_k, cache_v, k, v, **flash_args)
+
+
+class Attention(nn.Module):
+    """Self-Attention layer."""
+
+    def __init__(self, dim, num_heads, bias=True):
+        super(Attention, self).__init__()
+        self.qkv = nn.Linear(dim, dim * 3, bias=bias)
+        self.proj = nn.Linear(dim, dim, bias=bias)
+        self.head_dim = dim // num_heads
+        self.num_heads = num_heads
+        self.scale = self.head_dim**-0.5
+        self.cache = nn.Module()
+
+    def forward(self, x):
+        qkv_shape = (-1, x.size(1), 3, self.num_heads, self.head_dim)
+        q, k, v = self.qkv(x).view(qkv_shape).unbind(dim=2)
+        q, k = self.cache.forward_rotary(q, k, inplace=True)
+        o = self.cache.forward_flash(self, q, k, v)
+        return self.proj(o.flatten(2))
+
+
+class MLP(nn.Module):
+    """Two layers MLP."""
+
+    def __init__(self, dim, mlp_dim, bias=True):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(dim, mlp_dim, bias=bias)
+        self.fc2 = nn.Linear(mlp_dim, dim, bias=bias)
+        self.activation = nn.GELU()
+
+    def forward(self, x):
+        return self.fc2(self.activation(self.fc1(x)))
+
+
+class Block(nn.Module):
+    """Transformer block."""
+
+    def __init__(self, dim, num_heads, mlp_dim, bias=True):
+        super(Block, self).__init__()
+        self.attn = Attention(dim, num_heads, bias=bias)
+        self.mlp = MLP(dim, mlp_dim, bias=bias)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        x = self.attn(self.norm1(x)).add_(x)
+        return self.mlp(self.norm2(x)).add_(x)
+
+
+class Transformer(nn.Module):
+    """Causal transformer decoder."""
+
+    def __init__(self, depth, dim, num_heads, mlp_dim, vocab_size):
+        super(Transformer, self).__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.vocab_size = vocab_size
+        self.tok_embeddings = nn.Embedding(vocab_size, dim)
+        self.blocks = nn.ModuleList(Block(dim, num_heads, mlp_dim) for _ in range(depth))
+        self.norm = nn.LayerNorm(dim)
+        self.text_proj = nn.Linear(dim, vocab_size, bias=False)
+
+    def forward(self, prompts, tokens, start_pos=0):
+        prompt_len = prompts.size(1)
+        start_pos = start_pos + (prompt_len if start_pos > 0 else 0)
+        end_pos = start_pos + tokens.size(1) + (0 if start_pos > 0 else prompt_len)
+        self.cache.set_seq(start_pos, end_pos)
+        x = self.tok_embeddings(tokens)
+        x = x if start_pos > 0 else torch.cat([prompts, x], dim=1)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x[:, 0 if start_pos > 0 else prompt_len :])
+        return self.text_proj(x).float()
+
+
+class TextDecoder(nn.Module):
+    """Module to decode texts."""
+
+    def __init__(
+        self,
+        depth,
+        embed_dim,
+        num_heads,
+        mlp_ratio,
+        prompt_embed_dim,
+        max_seq_len,
+        vocab_size,
+    ):
+        super(TextDecoder, self).__init__()
+        self.max_seq_len = max_seq_len
+        self.max_text_len = self.max_seq_len - 1
+        self.encoder = nn.Linear(prompt_embed_dim, embed_dim, bias=False)
+        self.transformer = Transformer(
+            depth=depth,
+            dim=embed_dim,
+            mlp_dim=embed_dim * mlp_ratio,
+            num_heads=num_heads,
+            vocab_size=vocab_size,
+        )
+
+    def reset_cache(self, max_batch_size=1, max_seq_len=None):
+        device, dtype = self.encoder.weight.device, self.encoder.weight.dtype
+        max_seq_len = self.max_seq_len if max_seq_len is None else max_seq_len
+        num_heads, head_dim = self.transformer.num_heads, self.transformer.head_dim
+        self.transformer.cache = TransformerCache(device=device, dtype=dtype)
+        self.transformer.cache.init_seq(max_batch_size)
+        self.transformer.cache.init_rotary(max_seq_len, head_dim, theta=10000.0)
+        kv_cache_size = (max_batch_size, max_seq_len, num_heads, head_dim)
+        for blk in self.transformer.blocks:
+            blk.attn.__dict__["cache"] = self.transformer.cache
+            self.transformer.cache.init_kv(blk.attn, kv_cache_size) if not self.training else None
+
+    def get_prompts(self, prompt_tokens):
+        return self.encoder(prompt_tokens)
+
+    def get_outputs(self, inputs, start_pos=0):
+        return {"text_pred": self.transformer(inputs["prompts"], inputs["tokens"], start_pos)}
+
+    def forward(self, inputs, start_pos=0):
+        return self.get_outputs(inputs, start_pos)

tokenize_anything/modeling/text_tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

tokenize_anything/modeling/text_tokenizer.py

@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# Source from: https://github.com/facebookresearch/llama/blob/main/llama/tokenizer.py
+
+import os
+from logging import getLogger
+from typing import List
+
+from sentencepiece import SentencePieceProcessor
+
+
+logger = getLogger()
+
+
+class TextTokenizer:
+    """Tokenizing and encoding/decoding text using SentencePiece."""
+
+    def __init__(self, model_path=None):
+        """
+        Initializes the Tokenizer with a SentencePiece model.
+
+        Args:
+            model_path (str): The path to the SentencePiece model file.
+        """
+        if model_path is None:
+            model_path = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)), "text_tokenizer.model"
+            )
+        # reload tokenizer
+        assert os.path.isfile(model_path), model_path
+        self.sp_model = SentencePieceProcessor(model_file=model_path)
+        logger.info(f"Reloaded SentencePiece model from {model_path}")
+        # BOS / EOS token IDs
+        self.n_words: int = self.sp_model.vocab_size()
+        self.bos_id: int = self.sp_model.bos_id()
+        self.eos_id: int = self.sp_model.eos_id()
+        self.pad_id: int = self.sp_model.pad_id()
+        self.pad_id += self.n_words if self.pad_id < 0 else 0
+        logger.info(f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}")
+        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
+
+    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
+        """
+        Encodes a string into a list of token IDs.
+
+        Args:
+            s (str): The input string to be encoded.
+            bos (bool): Whether to prepend the beginning-of-sequence token.
+            eos (bool): Whether to append the end-of-sequence token.
+
+        Returns:
+            List[int]: A list of token IDs.
+        """
+        assert type(s) is str
+        t = self.sp_model.encode(s)
+        if bos:
+            t = [self.bos_id] + t
+        if eos:
+            t = t + [self.eos_id]
+        return t
+
+    def decode(self, t: List[int]) -> str:
+        """
+        Decodes a list of token IDs into a string.
+
+        Args:
+            t (List[int]): The list of token IDs to be decoded.
+
+        Returns:
+            str: The decoded string.
+        """
+        return self.sp_model.decode(t)
+
+    def tokenize(self, texts, context_length=None):
+        """Encode a list of string.
+
+        Parameters
+        ----------
+        texts : Union[str, List[str]]
+            The input text(s).
+        context_length : int, optional
+            The max token length.
+
+        Returns
+        -------
+        List[List[int]]
+            The encoded token indices.
+
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        tokens = [self.encode(text, bos=True, eos=True) for text in texts]
+        if context_length is None:
+            return tokens
+        truncated_tokens = []
+        for k, t in enumerate(tokens):
+            if len(t) > context_length:
+                t = t[:context_length]
+                t[-1] = self.eos_id
+            truncated_tokens.append(t)
+        return truncated_tokens
+
+    def detokenize(self, tokens):
+        """Decode a list of string.
+
+        Parameters
+        ----------
+        tokens : Union[List[List[int]], numpy.ndarray]
+            The input tokens.
+
+        Returns
+        -------
+        List[str]
+            The decoded text strings.
+
+        """
+        if hasattr(tokens, "tolist"):
+            tokens = tokens.tolist()
+        texts = []
+        for i in range(len(tokens)):
+            t = tokens[i][1:]
+            try:
+                eot_idx = t.index(self.eos_id)
+                t = t[:eot_idx]
+            except ValueError:
+                pass
+            texts.append(self.decode(t))
+        return texts

tokenize_anything/test_engine.py

@@ -0,0 +1,81 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Engine for testing."""
+
+import time
+
+from tokenize_anything.build_model import model_registry
+
+
+class InferenceCommand(object):
+    """Command to run batched inference."""
+
+    def __init__(self, input_queue, output_queue, kwargs):
+        self.input_queue = input_queue
+        self.output_queue = output_queue
+        self.kwargs = kwargs
+
+    def build_env(self):
+        """Build the environment."""
+        self.batch_size = self.kwargs.get("batch_size", 1)
+        self.batch_timeout = self.kwargs.get("batch_timeout", None)
+
+    def build_model(self):
+        """Build and return the model."""
+        builder = model_registry[self.kwargs["model_type"]]
+        return builder(device=self.kwargs["device"], checkpoint=self.kwargs["weights"])
+
+    def build_predictor(self, model):
+        """Build and return the predictor."""
+        return self.kwargs["predictor_type"](model, self.kwargs)
+
+    def send_results(self, predictor, indices, examples):
+        """Send the inference results."""
+        results = predictor.get_results(examples)
+        if hasattr(predictor, "timers"):
+            time_diffs = dict((k, v.average_time) for k, v in predictor.timers.items())
+            for i, outputs in enumerate(results):
+                self.output_queue.put((indices[i], time_diffs, outputs))
+        else:
+            for i, outputs in enumerate(results):
+                self.output_queue.put((indices[i], outputs))
+
+    def run(self):
+        """Main loop to make the inference outputs."""
+        self.build_env()
+        model = self.build_model()
+        predictor = self.build_predictor(model)
+        must_stop = False
+        while not must_stop:
+            indices, examples = [], []
+            deadline, timeout = None, None
+            for i in range(self.batch_size):
+                if self.batch_timeout and i == 1:
+                    deadline = time.monotonic() + self.batch_timeout
+                if self.batch_timeout and i >= 1:
+                    timeout = deadline - time.monotonic()
+                try:
+                    index, example = self.input_queue.get(timeout=timeout)
+                    if index < 0:
+                        must_stop = True
+                        break
+                    indices.append(index)
+                    examples.append(example)
+                except Exception:
+                    pass
+            if len(examples) == 0:
+                continue
+            self.send_results(predictor, indices, examples)

tokenize_anything/utils/__init__.py

@@ -0,0 +1,15 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------

tokenize_anything/utils/image.py

@@ -0,0 +1,73 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Image utilities."""
+
+import numpy as np
+import PIL.Image
+
+
+def im_resize(img, size=None, scale=None, mode="linear"):
+    """Resize image by the scale or size."""
+    if size is None:
+        if not isinstance(scale, (tuple, list)):
+            scale = (scale, scale)
+        h, w = img.shape[:2]
+        size = int(h * scale[0] + 0.5), int(w * scale[1] + 0.5)
+    else:
+        if not isinstance(size, (tuple, list)):
+            size = (size, size)
+    resize_modes = {"linear": PIL.Image.BILINEAR}
+    img = PIL.Image.fromarray(img)
+    return np.array(img.resize(size[::-1], resize_modes[mode]))
+
+
+def im_rescale(img, scales, max_size=0):
+    """Rescale image to match the detecting scales."""
+    im_shape = img.shape
+    img_list, img_scales = [], []
+    size_min = np.min(im_shape[:2])
+    size_max = np.max(im_shape[:2])
+    for target_size in scales:
+        im_scale = float(target_size) / float(size_min)
+        target_size_max = max_size if max_size > 0 else target_size
+        if np.round(im_scale * size_max) > target_size_max:
+            im_scale = float(target_size_max) / float(size_max)
+        img_list.append(im_resize(img, scale=im_scale))
+        img_scales.append((im_scale, im_scale))
+    return img_list, img_scales
+
+
+def im_vstack(arrays, fill_value=None, dtype=None, size=None, align=None):
+    """Stack image arrays in sequence vertically."""
+    if fill_value is None:
+        return np.vstack(arrays)
+    # Compute the max stack shape.
+    max_shape = np.max(np.stack([arr.shape for arr in arrays]), 0)
+    if size is not None and min(size) > 0:
+        max_shape[: len(size)] = size
+    if align is not None and min(align) > 0:
+        align_size = np.ceil(max_shape[: len(align)] / align)
+        max_shape[: len(align)] = align_size.astype("int64") * align
+    # Fill output with the given value.
+    output_dtype = dtype or arrays[0].dtype
+    output_shape = [len(arrays)] + list(max_shape)
+    output = np.empty(output_shape, output_dtype)
+    output[:] = fill_value
+    # Copy arrays.
+    for i, arr in enumerate(arrays):
+        copy_slices = (slice(0, d) for d in arr.shape)
+        output[(i,) + tuple(copy_slices)] = arr
+    return output

tokenize_anything/utils/mask.py

@@ -0,0 +1,45 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Mask utilities."""
+
+import numpy as np
+from pycocotools.mask import encode
+
+
+def mask_to_box(mask):
+    """Convert binary masks to boxes."""
+    shape, (h, w) = mask.shape, mask.shape[-2:]
+    masks = mask.reshape((-1, h, w)).astype("bool")
+    in_height = np.max(masks, axis=-1)
+    in_width = np.max(masks, axis=-2)
+    in_height_coords = in_height * np.arange(h, dtype="int32")
+    in_width_coords = in_width * np.arange(w, dtype="int32")
+    bottom_edges = np.max(in_height_coords, axis=-1)
+    top_edges = np.min(in_height_coords + h * (~in_height), axis=-1)
+    right_edges = np.max(in_width_coords, axis=-1)
+    left_edges = np.min(in_width_coords + w * (~in_width), axis=-1)
+    is_empty = (right_edges < left_edges) | (bottom_edges < top_edges)
+    boxes = np.stack([left_edges, top_edges, right_edges, bottom_edges], axis=-1)
+    boxes = boxes.astype("float32") * ((~is_empty)[:, None])
+    return boxes.reshape(*shape[:-2], 4) if len(shape) > 2 else boxes[0]
+
+
+def encode_masks(masks):
+    """Encode a set of masks to RLEs."""
+    rles = encode(np.asfortranarray(masks))
+    for rle in rles:
+        rle["counts"] = rle["counts"].decode()
+    return rles

tokenize_anything/utils/timer.py

@@ -0,0 +1,51 @@
+# ------------------------------------------------------------------------
+# Copyright (c) 2023-present, BAAI. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ------------------------------------------------------------------------
+"""Timing functions."""
+
+import contextlib
+import time
+
+
+class Timer(object):
+    """Simple timer."""
+
+    def __init__(self):
+        self.total_time = 0.0
+        self.calls = 0
+        self.start_time = 0.0
+        self.diff = 0.0
+        self.average_time = 0.0
+
+    def add_diff(self, diff, n=1, average=True):
+        self.total_time += diff
+        self.calls += n
+        self.average_time = self.total_time / self.calls
+        return self.average_time if average else self.diff
+
+    @contextlib.contextmanager
+    def tic_and_toc(self, n=1, average=True):
+        try:
+            yield self.tic()
+        finally:
+            self.toc(n, average)
+
+    def tic(self):
+        self.start_time = time.time()
+        return self
+
+    def toc(self, n=1, average=True):
+        self.diff = time.time() - self.start_time
+        return self.add_diff(self.diff, n, average)

tokenize_anything/version.py

@@ -0,0 +1,3 @@
+version = "0.1.0a0"
+git_version = "None"
+__version__ = version