Codes

.gitattributes

@@ -1,35 +1,36 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+flash_attn-2.6.1+cu118torch2.4cxx11abiFALSE-cp39-cp39-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,14 @@
 ---
-license: mit
+title: CalliDemo
+emoji: 📊
+colorFrom: red
+colorTo: pink
+sdk: gradio
+sdk_version: 5.15.0
+app_file: app.py
+pinned: false
+license: gpl
+short_description: A test demo of CalliReader
 ---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,21 @@
+import gradio as gr
+from PIL import Image
+from inference import single_image_wrapped
+
+# 调用 single_image_wrapped 函数的包装函数
+def image_with_prompt_to_result(image, prompt):
+    # 将图片传递给 single_image_wrapped 函数，得到输出结果
+    result = single_image_wrapped(image, prompt)
+    return result
+
+# 创建 Gradio 接口
+iface = gr.Interface(
+    fn=image_with_prompt_to_result,  # 调用的函数
+    inputs=[gr.Image(type="pil"), gr.Textbox(placeholder="输入提示词...")],  # 输入：图片 + prompt
+    outputs="text",  # 输出：识别结果（文本）
+    title="图片与提示词输入",  # 界面标题
+    description="上传一张图片并输入提示词，获取识别结果",  # 描述
+)
+
+# 启动应用
+iface.launch()

config/__init__.py

@@ -0,0 +1,5 @@
+import sys
+import os
+
+# 将folder1的路径添加到sys.path
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))

config/configu.py

@@ -0,0 +1,62 @@
+import torch
+import os
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 路径配置
+VIT_MODEL_PATH = '/home/luoyx/InternVL/CalliReader/params/vit_model.pt'
+MLP1_PATH = '/home/luoyx/InternVL/CalliReader/params/params/mlp1.pth'
+TOK_EMBEDDING_PATH = '/home/luoyx/InternVL/CalliReader/params/token_embedding.pth'
+TOKENIZER_PATH = 'InternVL'
+NORM_PARAMS_PATH='/home/luoyx/InternVL/CalliReader/params/gauss_norm_mu_sigma.pth'
+NORM_TOK_EMBEDDING_PATH='/home/luoyx/InternVL/CalliReader/params/gauss_norm.pth'
+NEW_1000_TOK_EMBEDDING_PATH='/home/luoyx/InternVL/CalliReader/params/new1000_token_embedding.pth'
+INTERNVL_PATH='InternVL'
+
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+SEED=42
+
+# 训练配置
+BATCH_SIZE = 256
+USE_WARMUP=False
+LR = 1e-4  # original 1e-4
+WEIGHT_DECAY = 1e-5
+WARMUP_STEPS = 2000 # *4 = total training steps
+NUM_EPOCHS = 13
+NUM_WORKERS = 4
+TRAIN_INTER = 10
+VAL_INTER = 500
+DOWNSAMPLE_RATIO = 0.5
+NUM_LAYERS=4
+GRAD_ACCU = 1
+MODEL_NAME = 'PERCEIVER'
+
+
+
+# 数据路径
+
+TRAIN_DATA_PATH = ""
+VAL_DATA_PATH = ''
+TEST_DATA_PATH = ''
+TRAIN_RATIO = 1#0.556 #0.02
+VAL_RATIO = 0.2#0.1
+
+# 36000 steps 8 cards 20 epochs, ~ 0.52 data ratio
+
+# LOGS andSAVE_NAME
+# 每一次跑新的实验切记一定需要修改！！！！
+LOG_NAME = ''
+SAVE_NAME = LOG_NAME+'.pth'
+
+# DDP
+WORLD_SIZE = torch.cuda.device_count()
+# 如果我们要加载训练一半的模型，两个都不能是none!!
+# LOAD CHECKPOINT AND RESUME TRAINING
+# PERCEIVER_CHECKPOINT ="/home/luoyx/InternVL/CalliReader/params/perceiver_4_n01_1e-4_new.pth"
+# RESUME = 26500
+PERCEIVER_CHECKPOINT ='/home/luoyx/InternVL/CalliReader/params/callialign.pth'
+RESUME = 50000
+ORDERFORMER_CHECKPOINT='/home/luoyx/InternVL/CalliReader/params/orderformer.pth'
+YOLO_CHECKPOINT="/home/luoyx/InternVL/CalliReader/params/best.pt"

flash_attn-2.6.1+cu118torch2.4cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

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

inference.py

@@ -0,0 +1,170 @@
+
+import random
+import numpy as np
+import torch
+from PIL import Image
+Image.MAX_IMAGE_PIXELS = None
+from transformers import AutoModel, AutoTokenizer
+import opencc
+from ultralytics import YOLO
+from config.configu import *
+from utils.utils import *
+import logging
+import argparse
+
+def setup_logger(log_file):
+    logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s - %(message)s')
+    logger = logging.getLogger()
+    return logger
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  
+   
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False 
+
+cc = opencc.OpenCC('t2s.json')
+set_seed(SEED)
+converter_t2s = opencc.OpenCC('t2s')
+
+
+def single_rec(model,tokenizer,detect_model,generation_config,image_path,prompt,use_p,hard_vq,drop_zero,repetition_penalty,verbose):
+    response, history = model.chat_ocr(tokenizer, detect_model,image_path, prompt, generation_config,
+                            use_p=use_p,
+                            hard_vq=hard_vq,
+                            drop_zero=drop_zero,repetition_penalty=repetition_penalty,return_history=True,verbose=verbose)
+    print(f'User: {prompt}\nAssistant: {response}')
+
+def folder_rec(model,tokenizer,detect_model,generation_config,folder_path,prompt,save_name,use_p,hard_vq,drop_zero,repetition_penalty,verbose):
+    results=[]
+
+    all_images=get_image_paths(folder_path)
+    for pic in tqdm(all_images):
+        pic_path=os.path.join(folder_path,pic)
+        try:
+            response, history = model.chat_ocr(tokenizer, detect_model,pic_path, prompt, generation_config,
+                                use_p=use_p,
+                                hard_vq=hard_vq,
+                                drop_zero=drop_zero,repetition_penalty=repetition_penalty,return_history=True,verbose=verbose)
+        except Exception as e:
+            print(f"An error has occured:\n{e}")
+            response="ERROR!"
+        print(f'User: {prompt}\nAssistant: {response}')
+        results.append({"imagePath":pic_path,'prompt':prompt,'response':response})
+    if not save_name.endswith('json'):
+        save_name+='_result.json'
+    save_json(save_name,results)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="args for inference task")
+
+    parser.add_argument('--tgt', type=str,help='Recognition target')
+    parser.add_argument('--prompt', type=str,default='这幅书法作品内容是什么？',help='Prompt for recognition')
+    parser.add_argument('--save_name',type=str,default="recognition.json",help="Storage of results if multiple images recognition mode")
+
+    parser.add_argument('--use_p', type=bool, default=True,help='Decide the usage of perceiver resampler')
+    parser.add_argument('--hard_vq', type=bool, default=False,help='Decide the usage of closest similarity match')
+    parser.add_argument('--drop_zero', type=bool, default=False,help='Decide the deletion of zero padding in pseudo tokens')
+    parser.add_argument('--verbose', type=bool, default=False,help='Decide the output of extra information')
+    parser.add_argument('--repetition_penalty', type=float, default=1.0,help='Repetition penalty for generation')
+    
+
+    args = parser.parse_args()
+
+    if not isinstance(args.tgt,str):
+        raise ValueError(f"The target should a string, not a instance of {type(args.tgt)}!")
+
+    
+    model = AutoModel.from_pretrained(
+            INTERNVL_PATH,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True).eval().cuda()
+    tokenizer = AutoTokenizer.from_pretrained(INTERNVL_PATH, trust_remote_code=True)
+
+    generation_config = dict(
+            num_beams=1,
+            max_new_tokens=1024,
+            do_sample=False,
+        )
+
+    detect_model=YOLO(YOLO_CHECKPOINT)
+    if is_image(args.tgt):
+        print("Single image recognition mode.")
+        single_rec(
+                    model,
+                    tokenizer,
+                    detect_model,
+                    generation_config,
+                    args.tgt,
+                    args.prompt,
+                    args.use_p,
+                    args.hard_vq,
+                    args.drop_zero,
+                    args.repetition_penalty,
+                    args.verbose)
+    elif os.path.isdir(args.tgt):
+        print("Multiple images recognition mode")
+        os.makedirs('results',exist_ok=True)
+        folder_rec(
+                   model,
+                   tokenizer,
+                   detect_model,
+                   generation_config,
+                   args.tgt,
+                   args.prompt,
+                   os.path.join('results',args.save_name),
+                   args.use_p,
+                   args.hard_vq,
+                   args.drop_zero,
+                   args.repetition_penalty,
+                   args.verbose)
+    else:
+        raise ValueError(f"The target should be either a image path or a folder that contain images!")
+
+def single_image_wrapped(image,prompts):
+    model = AutoModel.from_pretrained(
+            INTERNVL_PATH,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True).eval().cuda()
+    tokenizer = AutoTokenizer.from_pretrained(INTERNVL_PATH, trust_remote_code=True)
+
+    generation_config = dict(
+            num_beams=1,
+            max_new_tokens=1024,
+            do_sample=False,
+        )
+    detect_model=YOLO(YOLO_CHECKPOINT)
+
+
+    temp_dir = "temp_images"
+    os.makedirs(temp_dir, exist_ok=True)
+    
+    # 获取图片的文件名，并保存
+    temp_image_path = os.path.join(temp_dir, "uploaded_image.png")
+    image.save(temp_image_path)
+    single_rec(
+                    model,
+                    tokenizer,
+                    detect_model,
+                    generation_config,
+                    temp_image_path,
+                    prompts,
+                    True,
+                    False,
+                    True,
+                    1.2,
+                    False)
+if __name__=='__main__':
+
+    main()
+
+

models/__init__.py

@@ -0,0 +1,6 @@
+import sys
+import os
+
+# 将folder1的路径添加到sys.path
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))
+

models/model.py

@@ -0,0 +1,546 @@
+import torch
+import torch.nn as nn
+import sys
+import os
+import json
+
+from collections import OrderedDict
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+sys.path.append(project_root)
+
+from transformers import AutoConfig
+from InternVL.modeling_intern_vit import InternVisionModel
+from .perceiver_resampler import PerceiverResampler, MLP
+from config.configu import device, VIT_MODEL_PATH, MLP1_PATH, TOK_EMBEDDING_PATH, TOKENIZER_PATH,NORM_TOK_EMBEDDING_PATH,NORM_PARAMS_PATH
+
+
+def load_json(pth):
+    """加载json文件"""
+    with open(pth, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+    return data
+def load_vision_model(location='cpu'):
+    vit_config = AutoConfig.from_pretrained(TOKENIZER_PATH, trust_remote_code=True).vision_config
+    vision_model = InternVisionModel(vit_config).to(device).to(torch.bfloat16)
+    state_dict = torch.load(VIT_MODEL_PATH, weights_only=True, map_location=location)
+    incompatible_keys = vision_model.load_state_dict(state_dict)
+    if incompatible_keys.unexpected_keys:
+        print(f"Unexpected keys: {incompatible_keys.unexpected_keys}")
+    if incompatible_keys.missing_keys:
+        print(f"Missing keys: {incompatible_keys.missing_keys}")
+    print("vision model已加载")
+    return vision_model
+
+def load_mlp1(downsample_ratio, vit_hidden_size=1024, llm_hidden_size=4096,location='cpu'):
+    mlp1 = nn.Sequential(
+        nn.LayerNorm(vit_hidden_size * int(1 / downsample_ratio) ** 2),
+        nn.Linear(vit_hidden_size * int(1 / downsample_ratio) ** 2, llm_hidden_size),
+        nn.GELU(),
+        nn.Linear(llm_hidden_size, llm_hidden_size)
+    ).to(device).to(torch.bfloat16)
+    mlp1.load_state_dict(torch.load(MLP1_PATH, weights_only=True, map_location=location))
+    print("mlp1已加载")
+    return mlp1
+
+def load_tok_embeddings(path=TOK_EMBEDDING_PATH,vocab_size=92553, llm_hidden_size=4096,location='cpu'):
+    tok_embeddings = nn.Embedding(vocab_size, llm_hidden_size, padding_idx=2).to(device).to(torch.bfloat16)
+    tok_embeddings.load_state_dict(torch.load(path, weights_only=True, map_location=location))
+    print("tok_embedding已加载")
+    return tok_embeddings
+
+
+
+def load_normed_tok_embeddings(vocab_size=92553, llm_hidden_size=4096,load_checkboard=False,location="cpu"):
+    tok_embeddings = nn.Embedding(vocab_size, llm_hidden_size, padding_idx=2).to(device).to(torch.bfloat16)
+    tok_embeddings.load_state_dict(torch.load(NORM_TOK_EMBEDDING_PATH, weights_only=True, map_location=location))
+    print("norm tok_embedding已加载")
+    if load_checkboard:
+        checkboard_norm=torch.load(NORM_PARAMS_PATH) # (voc_size, 2) mu sigma    pred * sigma + mu (逐行)
+        print("归一化参数(mu, sigma)已加载")
+        return tok_embeddings,checkboard_norm
+    return tok_embeddings
+
+
+def load_tokenizer():
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
+    return tokenizer
+
+def load_perceiver_resampler(path=None, num_layers=4, checkpoint=None):
+    model = PerceiverResampler(dim=4096, depth = num_layers).to(device).to(torch.bfloat16)
+    if checkpoint == None and path!=None:
+        checkpoint = torch.load(path)
+    if path is not None:
+        print(f"Load from {path}")
+        if isinstance(checkpoint, dict):
+            if 'model_state_dict' in checkpoint.keys():
+                model.load_state_dict(checkpoint['model_state_dict'])
+            else:
+                raise FileNotFoundError("no key model_state_dict in ckpt")
+        else:
+            model.load_state_dict(checkpoint)
+    print(f"Model has a parameter scale of {sum(p.numel() for p in model.parameters())/1e9:.3f} B.")
+    return model
+
+def load_mlp(path=None):
+    model = MLP(dim=256).to(device).to(torch.bfloat16)
+    if path is not None:
+        model.load_state_dict(torch.load(path))
+    print(f"Model has a parameter scale of {sum(p.numel() for p in model.parameters())/1e9:.3f} B.")
+    return model
+
+def load_perceiver_resampler_2(model_path, num_layers=4,device=None):
+    if device==None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # 初始化模型
+    model = PerceiverResampler(dim=4096,depth=num_layers)
+    
+    # 加载预训练权重
+    state_dict = torch.load(model_path, map_location='cpu')
+    
+    # 移除 state_dict 中的 `module.` 前缀
+    state_dict = torch.load(model_path,weights_only=False)
+    if 'model_state_dict' in state_dict.keys():
+        state_dict = state_dict['model_state_dict']
+
+# 3. 处理 DDP 模型的情况（检查是否有 'module.' 前缀）
+    new_state_dict = OrderedDict()
+
+    for key, value in state_dict.items():
+        # 如果有 'module.' 前缀，则去掉它
+        if key.startswith('module.'):
+            new_key = key[len('module.'):]
+        else:
+            new_key = key
+        new_state_dict[new_key] = value
+    model = model.to_empty(device=device)
+    model.load_state_dict(new_state_dict)
+    
+    # 将模型移动到目标设备
+    
+    
+    # 将模型转换为所需的数据类型
+    model = model.to(torch.bfloat16)
+    return model
+
+def load_pretrained_resampler(checkpoint_path, num_layers=6):
+    model = load_perceiver_resampler(num_layers=num_layers)
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    #print(checkpoint.keys())
+    # 如果模型是通过 DDP 保存的，需要处理 'module.' 前缀
+    if 'module.' in list(checkpoint.keys())[0]:
+        print("load ddp Perceiver Resampler....")
+        #model = torch.nn.parallel.DistributedDataParallel(model)
+        model.load_state_dict(checkpoint)
+    elif 'module.' in list(checkpoint['model'].keys())[0]:
+        print("load ddp Perceiver Resampler....")
+        #model = torch.nn.parallel.DistributedDataParallel(model)
+        model.load_state_dict(checkpoint['model'])
+    else:
+        print("load Perseiver Resampler ...")
+        model.load_state_dict(checkpoint)
+    return model
+
+def load_optimizer(optimizer, path, resume):
+    # 加载checkpoint
+    
+    if resume:
+        #assert isinstance(ckpt, dict) and 'optimizer_state_dict' in ckpt
+        ckpt = torch.load(path)
+        if 'optimizer_state_dict' not in ckpt:
+            return optimizer
+        # 处理 DDP 模型的情况
+        optimizer_state_dict = ckpt['optimizer_state_dict']
+        new_optimizer_state_dict = OrderedDict()
+
+        for key, value in optimizer_state_dict.items():
+            if key.startswith('module.'):
+                new_key = key[len('module.'):]
+            else:
+                new_key = key
+            new_optimizer_state_dict[new_key] = value
+
+        optimizer.load_state_dict(new_optimizer_state_dict)
+    
+    return optimizer
+
+def load_scheduler(scheduler, path, resume):
+    
+    if resume:
+        #assert isinstance(ckpt, dict) and 'scheduler_state_dict' in ckpt
+        # 加载checkpoint
+        ckpt = torch.load(path)
+        if 'scheduler_state_dict' not in ckpt:
+            return scheduler
+        # 处理 DDP 模型的情况
+        scheduler_state_dict = ckpt['scheduler_state_dict']
+        new_scheduler_state_dict = OrderedDict()
+        for key, value in scheduler_state_dict.items():
+            if key.startswith('module.'):
+                new_key = key[len('module.'):]
+            else:
+                new_key = key
+            new_scheduler_state_dict[new_key] = value
+
+        scheduler.load_state_dict(new_scheduler_state_dict)
+    
+    return scheduler
+
+import numpy as np
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset, random_split
+class BoundingBoxDataset(Dataset):
+    """数据集class"""
+    def __init__(self, data, targets):
+        self.data = data
+        self.targets = targets
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        x = self.data[idx]
+        y = self.targets[idx]
+        return x, y
+
+class Transformer(nn.Module):
+    """核心的Transformer model,encoder only"""
+    def __init__(self, input_dim:int, model_dim:int, num_heads:int, num_layers:int,output_dim:int,norms=True):
+        super(Transformer, self).__init__()
+        self.embedding=nn.Linear(input_dim,model_dim)
+        if norms:
+            self.layer_norm = nn.LayerNorm(model_dim)
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads,batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers,norm=self.layer_norm if norms==True else None)
+    
+        self.decoder=nn.Linear(model_dim,output_dim)
+
+    def forward(self, x):
+        x=self.embedding(x)
+        x = self.transformer_encoder(x)
+        x=self.decoder(x)
+        return x
+    
+class OrderFormer:
+    """封装后的模型,实现数据加载,训练,测试,推理功能"""
+    def __init__(self, model_path=None,max_nums=300,input_dim=4, model_dim=256, num_heads=8, num_layers=4, output_dim=1,device=torch.device("cuda"),label_name="turn",norm=False):
+        self.model = Transformer(input_dim, model_dim, num_heads, num_layers, output_dim,norms=norm).to_empty(device=device)
+        if isinstance(model_path,str):
+            self.model.load_state_dict(torch.load(model_path))
+
+        self.device=device
+        self.max_nums=max_nums
+        self.input_dim=input_dim
+        self.label_name=label_name
+
+    def _get_all_jsons(self,folder_path):
+        """得到文件夹中的所有json文件路径"""
+        files = os.listdir(folder_path)
+        json_files = [folder_path+f for f in files if os.path.isfile(os.path.join(folder_path, f)) and f.endswith('json')]
+        return json_files
+    
+    def _preprocess(self,datas):
+        """
+        data: SHOULD BE Consistent with labelme data format
+        return: 
+            [
+                [
+                    [x1,y1,x2,y2],label
+                ]
+                ...
+            ]
+        x,y:[0,1]    
+        """
+        data=datas['shapes']
+        h=datas['imageHeight']
+        w=datas['imageWidth']
+        example=[]
+        X=[]
+        Y=[]
+        L=[]
+        for obj in data:
+                #记录顺序,横纵坐标
+            l=obj[self.label_name]
+            p=obj['points']
+            X.extend([p[0][0]/w,p[1][0]/w])
+            Y.extend([p[0][1]/h,p[1][1]/h])
+            L.append(l)
+        xmin=min(X)
+        ymin=min(Y)
+            #横纵坐标均减去最小值,保持平移不变性
+        X=np.array(X)-xmin
+        Y=np.array(Y)-ymin
+        for i in range(len(L)):
+            coord=[X[2*i],Y[2*i],X[2*i+1],Y[2*i+1]]
+            example.append([coord,L[i]])
+        return example
+    def _sort_boxes(self,boxes):
+        """以到(0,0)距离排序box,确保输入box是唯一的排列序列
+        boxes=[[[x1,y1,x2,y2],label],...]
+        label可以是标签,也可以是原始的bbox便于得到bbox和顺序的对应关系
+        """
+        return sorted(boxes,key=lambda x:((x[0][0]+x[0][2])/2)**2+((x[0][1]+x[0][3])/2)**2)
+
+    def _load_data(self,path,device=torch.device("cuda"),name='turn'):
+        """
+        从json转为tensor的构造函数
+        Args:
+        path:jsons-jpgs所存在的文件夹
+        max_nums:单个样本中char的最大个数
+        name:取得char顺序指标的key
+        Return:
+
+        """
+        max_nums=self.max_nums
+        device=self.device
+        all_jsons=self._get_all_jsons(path)
+        raw=[]
+        for j in all_jsons:
+            datas=load_json(j)
+            example=self._preprocess(datas)
+            raw.append(example)
+        transformed_inputs=[]
+        transformed_labels=[]
+        originNs=[]#记录原序列的长度,用于从结果中得到序列
+        for item in raw:
+            item=self._sort_boxes(item)
+            originNs.append(len(item))
+            lst=[]
+            ls=[]
+            for x in item:
+                #lst=lst+[x1,y1,x2,y2]
+                lst.extend(x[0])
+                #ls记录label
+                ls.append(int(x[1]))
+            #pad全0序列和全0标签到指定的max_nums长度
+            lst.extend([0]*self.input_dim*(max_nums-len(item)))
+            ls.extend([0]*(max_nums-len(item)))
+            
+            transformed_inputs.append(lst)
+            transformed_labels.append(ls)
+        return torch.tensor(transformed_inputs,dtype=torch.float32).reshape((-1,max_nums,self.input_dim)).to(device),torch.tensor(transformed_labels,dtype=torch.float32).reshape((-1,self.max_nums,1)).to(device),originNs
+    
+    def _decode(self,output,N,batch_size=1):
+        """从输出的tensor解码得到排序"""
+        new_output=output.reshape((batch_size,-1))[:,:N]
+        sorted_indices = torch.argsort(new_output, dim=1)
+        ranks = torch.argsort(sorted_indices, dim=1)
+        return ranks + 1
+    
+    def _get_acc(self,tensor1, tensor2):
+        """计算两个相同形状tensor数值相同的位置的占比"""
+        # Ensure the tensors are of the same shape
+        assert tensor1.shape == tensor2.shape, "Tensors must have the same shape"
+        
+        # Create a boolean mask where the values are equal
+        equal_mask = tensor1 == tensor2
+        
+        # Calculate the proportion of equal values
+        equal_count = torch.sum(equal_mask).item()
+        total_elements = torch.numel(tensor1)
+        
+        proportion_equal = equal_count / total_elements
+        
+        return proportion_equal
+
+
+    def train(self, path,batch_size=4,lr=0.0002,weight_decay=0,epochs=1000,verbose=True):
+        """训练函数"""
+        if verbose:
+            print("Loading dataset...")
+        data,labels,_=self._load_data(path=path,device=self.device,name=self.label_name)
+
+        #TODO :可指定的训练策略
+        optimizer = optim.AdamW(self.model.parameters(), lr=lr,weight_decay=weight_decay,amsgrad=True)
+        # scheduler=optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=10)
+        scheduler=optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=1e-6)
+        criterion=torch.nn.MSELoss()
+
+        dataset = BoundingBoxDataset(data, labels)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+        min_loss=float("inf")
+        if verbose:
+            print("Start training...")
+        for epoch in range(epochs):
+            losses=0
+            for batch_idx,(inputs, y) in enumerate(tqdm((dataloader))):                     
+                optimizer.zero_grad()
+                outputs = self.model(inputs)  
+
+                loss = criterion(outputs, y)
+                loss.backward()
+                losses+=loss.item()
+                scheduler.step(epoch + batch_idx / len(dataloader))
+                optimizer.step()
+            #scheduler.step()
+            
+            if verbose:
+                print(f"Epoch {epoch+1}/{epochs}, Loss: {losses/len(dataloader)}")
+            if losses/len(dataloader)<min_loss:
+                min_loss=losses/len(dataloader)
+                if verbose:
+                    print("Saving best model...")
+                torch.save(self.model.state_dict(),'best.pth')
+
+ 
+    def eval(self, path,verbose=False):
+        """在数据集上测试,计算平均loss和mAP"""
+        testdata,testlabels,Ns=self._load_data(path=path,device=self.device,name=self.label_name)
+        dataset = BoundingBoxDataset(testdata, testlabels)
+        testloader=DataLoader(dataset,batch_size=1,shuffle=False)
+
+        self.model.eval()
+        losses=0
+        mAP=0
+        if verbose:
+            print("Evaluation...")
+        criterion = nn.MSELoss()
+        for i,(inputs, y) in enumerate(testloader):
+            outputs = self.model(inputs)
+            pred= self._decode(outputs,Ns[i])
+            gt=y.reshape((1,-1))[:,:Ns[i]]
+            loss = criterion(pred, gt)
+            acc=self._get_acc(pred,gt)
+            if verbose:
+
+                print("Pred:",pred)
+                print("GT:",gt)
+                print("loss= ",loss.item())
+                print("acc= ",acc,'\n')
+            losses+=loss.item()
+            #mAP+=1 if acc==1 else 0
+            mAP+=acc
+        print(f"Test MSELoss= {losses/len(testloader):.4f}\nTest mAP= {mAP/len(testloader):.4f}")
+    
+    def predict(self,datas,jpg_path=None,save_path=None,verbose=False):
+        """
+        进行单个数据的预测,如果有图片,保存路径,可以进行verbose可视化
+        返回一个dict,key是顺序,value是box的位置
+        """
+        if save_path:
+            os.makedirs(save_path,exist_ok=True)
+        import time
+        st=time.time()
+        data=datas['shapes']
+        h=datas['imageHeight']
+        w=datas['imageWidth']
+        example=[]
+        X=[]
+        Y=[]
+        Ls=[]
+        for obj in data:
+                #记录顺序,横纵坐标
+            p=obj['points']
+            flat_p=[p[0][0],p[0][1],p[1][0],p[1][1]]
+            Ls.append(flat_p)
+            X.extend([p[0][0]/w,p[1][0]/w])
+            Y.extend([p[0][1]/h,p[1][1]/h])
+        xmin=min(X)
+        ymin=min(Y)
+            #横纵坐标均减去最小值,保持平移不变性
+        X=np.array(X)-xmin
+        Y=np.array(Y)-ymin
+        for i in range(len(data)):
+            coord=[X[2*i],Y[2*i],X[2*i+1],Y[2*i+1]]
+            example.append([coord,Ls[i]])
+        example=self._sort_boxes(example)
+        inputs=[]
+        labels=[]
+        for coord in example:
+            inputs.extend(coord[0])
+            labels.append(coord[1])
+        inputs.extend([0]*self.input_dim*(self.max_nums-len(example)))
+        
+        x=torch.tensor(inputs,dtype=torch.bfloat16).reshape((-1,self.max_nums,self.input_dim)).to(self.device)
+        
+        mstart=time.time()
+        self.model.eval()
+        y=self.model(x)
+        mtime=time.time()-mstart
+        pred=self._decode(y,len(example)).squeeze().tolist()
+        results={}
+        if isinstance(pred,int):
+            pred=[pred]
+        for p,l in zip(pred,labels):
+            results[p]=l
+        
+        post_start=time.time()
+        results=self.postprocess(dict(sorted(results.items(), key=lambda item: item[0])),w,h,save_path,jpg_path)
+        ptime=time.time()-post_start
+        if verbose:
+            print(f"Using {time.time()-st:.3f}s to sort boxes,with {mtime:.3f}s on OrderFormer inference,{ptime:.3f}s on postprocess.")
+        if verbose and isinstance(jpg_path,str) and isinstance(save_path,str):
+            import cv2
+            frame = cv2.imread(jpg_path)
+
+            for idx ,points in results.items():
+                x1, y1, x2, y2 = int(points[0]), int(points[1]), int(points[2]), int(points[3])
+                cv2.rectangle(frame, (x1, y1), (x2, y2), thickness=2,color=(255,0,0),lineType=cv2.LINE_AA)
+                label_position = ((x1+x2)//2,(y1+y2)//2)  # Adjust the position of the label as needed
+                cv2.putText(frame, str(idx), label_position, cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, cv2.LINE_AA)
+            name=jpg_path.split("/")[-1]
+            cv2.imwrite(save_path+"ordered_"+name,frame)
+    
+        return dict(sorted(results.items(), key=lambda item: item[0]))
+
+    
+    
+    def postprocess(self,results,width,height,save_dir,jpg_path,vis=True,max_iters=5):
+        def ordered_permute(b1,b2,b3):
+            ws=[b1[2]-b1[0],b2[2]-b2[0],b3[2]-b3[0]]
+            hs=[b1[3]-b1[1],b2[3]-b2[1],b3[3]-b3[1]]
+            c1=[(b1[0]+b1[2])/2,(b1[1]+b1[3])/2]
+            c2=[(b2[0]+b2[2])/2,(b2[1]+b2[3])/2]
+            c3=[(b3[0]+b3[2])/2,(b3[1]+b3[3])/2]
+            s=[ws[0]*hs[0],ws[1]*hs[1],ws[2]*hs[2]]
+            if max(abs(c1[1]-c2[1]),abs(c1[1]-c3[1]),abs(c2[1]-c3[1]))<min(hs) and min(s)/max(s)>0.7:
+                c=[c1[0],c2[0],c3[0]]
+
+            else:
+                c=[3,2,1]
+            indexed_c = list(enumerate(c))  
+
+            
+            sorted_by_value = sorted(indexed_c, key=lambda x: x[1],reverse=True)  
+
+      
+            sorted_indices = [index for index, value in sorted_by_value]
+
+            return sorted_indices
+        index=list(results.keys())
+        boxes=[[item[0]/width,item[1]/height,item[2]/width,item[3]/height] for item in list(results.values())]
+        for i in range(len(index)-2):
+            now=boxes[i]
+            next_1=boxes[i+1]
+            next_2=boxes[i+2]
+            order=ordered_permute(now,next_1,next_2)
+            
+            j=i+1
+            boxes[i],boxes[i+1],boxes[i+2]=boxes[i+order[0]],boxes[i+order[1]],boxes[i+order[2]]
+            results[j],results[j+1],results[j+2]=results[j+order[0]],results[j+order[1]],results[j+order[2]]
+            
+        return results
+    
+def load_orderformer(path,
+    max_num=50,
+    input_dim=4,
+    output_dim=1,
+    model_dim=256,
+    num_layers=4,
+    num_heads=8,
+    ):
+
+    model=OrderFormer(max_nums=max_num,
+                        num_layers=num_layers,
+                        input_dim=input_dim,
+                        output_dim=output_dim,
+                        model_dim=model_dim,
+                        num_heads=num_heads,
+                        model_path=path,
+                        label_name='turn',
+                        norm=False)
+    return model

models/perceiver_resampler.py

@@ -0,0 +1,154 @@
+import os
+import torch
+from torch import nn
+from einops import rearrange, repeat
+from torch import einsum
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=64,
+        heads=8
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        inner_dim = dim_head * heads  # 512
+
+        self.norm_media = nn.LayerNorm(dim)
+        self.norm_learns = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False) # 4096×512
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) # 4096×1024
+        self.to_out = nn.Linear(inner_dim, dim, bias=False) # 512×4096
+
+    def forward(self, x, learns): # x(b, 256, 4096), learns(b, 3, 4096)
+        x = self.norm_media(x)
+        learns = self.norm_learns(learns)
+
+        b, n, h = *x.shape[:2], self.heads
+
+        q = self.to_q(learns) # q(b, 3, 512)
+
+        # 注意：在PerceiverResampler中，将输入和learns拼接后进行attention计算
+        kv_input = torch.cat((x, learns), dim=-2) # kv_input(b, 259, 4096)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1) # (b, 259, 1024)->k, v(b, 259, 512)
+
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v)) # q(b, 8, 3, 64)   k, v(b, 8, 259, 64)
+
+        q = q * self.scale
+
+        # attention计算
+        sim = einsum('b h i d, b h j d -> b h i j', q, k)
+        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
+        attn = sim.softmax(dim=-1) # sim, attn(b, 8, 3, 259)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v) # out(b, 8, 3, 64)
+        out = rearrange(out, 'b h n d -> b n (h d)') # out(b, 3, 512)
+        return self.to_out(out) # return(b, 3, 4096)
+
+
+class PerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,  
+        depth=6,
+        dim_head=64,
+        heads=8,
+        num_learns=3,  
+        ff_mult=4,
+    ):
+        super().__init__()
+        self.learns = nn.Parameter(torch.randn(num_learns, dim))
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, 256, 4096)
+        Returns:
+            shape (b, 3, 4096) where 3 is self.num_learns
+        """
+        b, n, d = x.shape 
+
+ 
+        learns = repeat(self.learns, "n d -> b n d", b=b)
+
+       
+        for attn, ff in self.layers:
+            
+            learns = attn(x, learns) + learns
+            learns = ff(learns) + learns
+
+        return self.norm(learns)
+    
+class MLP(nn.Module):
+    def __init__(self, input_dim, hidden_mult=4): # input_dim = 256
+        super().__init__()
+        self.ff1 = FeedForward_2(input_dim, input_dim, hidden_mult)
+        self.ff2 = FeedForward_2(input_dim, 3, hidden_mult)
+
+    def forward(self, x):
+  
+        x = x.permute(0, 2, 1)
+        x = self.ff1(x)
+        x = self.ff2(x)
+        
+        x = x.permute(0, 2, 1)
+        return x
+    
+class MLP_6763(nn.Module):
+    def __init__(self, input_dim, output_dim, hidden_mult=2): 
+        super().__init__()
+        self.ff1 = FeedForward_2(input_dim, output_dim, hidden_mult)
+        self.ff2 = FeedForward_2(output_dim, output_dim, hidden_mult)
+
+    def forward(self, x):
+        b, n, d = x.shape
+        x = x.view(b, -1)
+        x = self.ff1(x)
+        x = self.ff2(x)
+        return x
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim * mult),
+            nn.GELU(),
+            nn.Linear(dim * mult, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+    
+class FeedForward_2(nn.Module):
+    def __init__(self, input_dim, output_dim, mult=4):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(input_dim),
+            nn.Linear(input_dim, input_dim * mult),
+            nn.GELU(),
+            nn.Linear(input_dim * mult, output_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)

models/similarity.py

@@ -0,0 +1,53 @@
+
+
+from __future__ import print_function
+import torch
+import torch.nn as nn
+from .model import*
+import torch.nn.functional as F
+
+def vq_cos_sim(embedding, input_tensor, use_dynamic_p=False,ddp=False):
+
+    if ddp:
+        embedding_weight = embedding.module.weight
+    else:
+        embedding_weight = embedding.weight  
+
+    
+    input_norm = F.normalize(input_tensor, p=2, dim=2)  
+    embedding_norm = F.normalize(embedding_weight, p=2, dim=1) 
+
+    similarity = torch.matmul(input_norm, embedding_norm.t())  
+    cos_sim_values, indices = similarity.max(dim=2)
+
+    if use_dynamic_p:
+        
+        return indices.squeeze(), cos_sim_values.squeeze()
+    
+    return indices.squeeze()
+
+
+class RatioLossWithMSELoss(nn.Module):
+    def __init__(self, total_iters, min_weight=0.001, max_weight=1,eps=torch.tensor(1e-3, dtype=torch.bfloat16)):
+        super(RatioLossWithMSELoss, self).__init__()
+        self.eps = eps
+        self.total_iters = total_iters
+        self.min_weight = min_weight
+        self.max_weight = max_weight
+        self.mse=nn.MSELoss()
+        self.losses={}
+    def forward(self, output, target, current_iter):
+
+        weight = self.min_weight + (self.max_weight - self.min_weight) * (current_iter / self.total_iters)
+        loss = (torch.abs(target - output)) / (torch.abs(target) + self.eps)
+        weighted_loss = weight * loss
+
+        self.losses['ratio']=loss.mean()
+        self.losses['mse']=self.mse(output,target)
+        return weighted_loss.mean()+self.mse(output,target)
+    
+
+
+
+if __name__=='__main__':
+    pass

requirements.txt

@@ -0,0 +1,86 @@
+accelerate==1.3.0
+aiohappyeyeballs
+aiohttp==3.10.5
+aiosignal==1.2.0
+async-timeout
+attrs==23.1.0
+Bottleneck
+Brotli 
+certifi 
+charset-normalizer 
+contourpy
+cycler
+decord==0.6.0
+dill 
+einops==0.8.0
+filelock 
+
+fonttools==4.51.0
+frozenlist==1.4.0
+fsspec 
+gmpy2 
+huggingface_hub 
+idna
+importlib-metadata==7.0.1
+importlib_resources
+jieba 
+Jinja2
+joblib==1.4.2
+kiwisolver 
+Levenshtein==0.26.1
+MarkupSafe 
+matplotlib==3.9.2
+mkl-service==2.4.0
+mkl_fft 
+mkl_random 
+mpmath 
+multidict==6.0.4
+multiprocess==0.70.15
+networkx 
+numexpr 
+numpy 
+OpenCC==1.1.6
+opencv-python
+opencv-python-headless
+packaging 
+pandas
+pillow
+propcache==0.2.1
+psutil 
+py-cpuinfo
+pyarrow 
+pyarrow-hotfix 
+pyparsing
+PySocks 
+python-dateutil 
+pytz 
+PyYAML 
+RapidFuzz==3.11.0
+regex==2024.11.6
+requests 
+rouge 
+safetensors==0.5.2
+scikit-learn==1.5.1
+scipy
+seaborn
+sentencepiece==0.2.0
+six 
+sympy 
+threadpoolctl==3.5.0
+timm==0.9.12
+tokenizers==0.20.3
+torch==2.4.0
+torchaudio==2.4.0
+torchvision==0.19.0
+tqdm 
+transformers==4.45.2
+triton==3.0.0
+typing_extensions 
+tzdata 
+ultralytics
+unicodedata2
+urllib3
+xxhash 
+yarl==1.11.0
+zipp==3.17.0
+./flash_attn-2.6.1+cu118torch2.4cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

test.py

@@ -0,0 +1 @@
+from InternVL import modeling_internvl_chat