inference_15.py

import torch

import re
import json


from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor


def standardize_messages(messages):
    for message in messages:
        if isinstance(message["content"], str):
            # 如果是字符串，转换成标准格式
            message["content"] = [
                {
                    "type": "text",
                    "text": message["content"]
                }
            ]
    return messages


def add_box_token(input_string):
    # Step 1: Split the string into individual actions
    if "Action: " in input_string and "start_box=" in input_string:
        suffix = input_string.split("Action: ")[0] + "Action: "
        actions = input_string.split("Action: ")[1:]
        processed_actions = []
        for action in actions:
            action = action.strip()
            # Step 2: Extract coordinates (start_box or end_box) using regex
            coordinates = re.findall(r"(start_box|end_box)='\((\d+),\s*(\d+)\)'", action)
            
            updated_action = action  # Start with the original action
            for coord_type, x, y in coordinates:
                # Convert x and y to integers
                updated_action = updated_action.replace(f"{coord_type}='({x},{y})'", f"{coord_type}='<|box_start|>({x},{y})<|box_end|>'")
            processed_actions.append(updated_action)
        
        # Step 5: Reconstruct the final string
        final_string = suffix + "\n\n".join(processed_actions)
    else:
        final_string = input_string
    return final_string


def standardize_messages(messages):
    for message in messages:
        if isinstance(message["content"], str):
            # 如果是字符串，转换成标准格式
            message["content"] = [
                {
                    "type": "text",
                    "text": message["content"]
                }
            ]
    return messages


result = {}
messages = json.load(open("./data/test_messages_07.json"))

for message in messages:
    if message["role"] == "assistant":
        message["content"] = add_box_token(message["content"])
        print(message["content"])


messages = standardize_messages(messages)


if __name__ == "__main__":
    from argparse import ArgumentParser
    
    from utils.coordinate_extract import extract_coordinates
    from utils.plot_image import plot
    
    parser = ArgumentParser()
    
    parser.add_argument("--model_path", type=str, required=True)
    
    args = parser.parse_args()
    
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            args.model_path,
            torch_dtype=torch.float16,
            device_map="auto",
            attn_implementation="sdpa"
        )
    
    processor = AutoProcessor.from_pretrained(args.model_path)
    
    
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to("cuda")
    
    
    generated_ids = model.generate(**inputs, max_new_tokens=400)
    
    generated_ids_trimmed = [
                out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    
    print(output_text)
    
    model_raw_response = output_text[0]
    
    coordinate = extract_coordinates(model_raw_response)

    # Please use re to parse the coordinate values
    model_output_width, model_output_height = coordinate
    
    plot(model_output_width, model_output_height)