File size: 9,306 Bytes
d2468c2 183355e d2468c2 cd3a7b0 d2468c2 e73817e d2468c2 945a97b d2468c2 0d409d1 d2468c2 5367a16 d2468c2 945a97b d2468c2 945a97b d2468c2 f9d46da d2468c2 f9d46da d2468c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 |
import argparse
import os
import re
import torch
from PIL import Image, ImageDraw
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from typing import List
import json
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
class AITM_Dataset(Dataset):
def __init__(self, json_path):
#self.data = []
with open(json_path, 'r') as f:
self.data = json.load(f)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
x = self.data[idx]
img_path = x['image']
task = x['conversations'][0]['value']
return img_path, task
def draw_boxes_on_image(image: Image.Image, boxes: List[List[float]], save_path: str):
"""
Draws red bounding boxes on the given image and saves it.
Parameters:
- image (PIL.Image.Image): The image on which to draw the bounding boxes.
- boxes (List[List[float]]): A list of bounding boxes, each defined as [x_min, y_min, x_max, y_max].
Coordinates are expected to be normalized (0 to 1).
- save_path (str): The path to save the updated image.
Description:
Each box coordinate is a fraction of the image dimension. This function converts them to actual pixel
coordinates and draws a red rectangle to mark the area. The annotated image is then saved to the specified path.
"""
draw = ImageDraw.Draw(image)
for box in boxes:
x_min = int(box[0] * image.width)
y_min = int(box[1] * image.height)
x_max = int(box[2] * image.width)
y_max = int(box[3] * image.height)
draw.rectangle([x_min, y_min, x_max, y_max], outline="red", width=3)
image.save(save_path)
def main():
"""
A continuous interactive demo using the CogAgent1.5 model with selectable format prompts.
The output_image_path is interpreted as a directory. For each round of interaction,
the annotated image will be saved in the directory with the filename:
{original_image_name_without_extension}_{round_number}.png
Example:
python cli_demo.py --model_dir THUDM/cogagent-9b-20241220 --platform "Mac" --max_length 4096 --top_k 1 \
--output_image_path ./results --format_key status_action_op_sensitive
"""
parser = argparse.ArgumentParser(
description="Continuous interactive demo with CogAgent model and selectable format."
)
parser.add_argument(
"--model_dir", required=True, help="Path or identifier of the model."
)
parser.add_argument(
"--platform",
default="Mac",
help="Platform information string (e.g., 'Mac', 'WIN').",
)
parser.add_argument(
"--max_length", type=int, default=4096, help="Maximum generation length."
)
parser.add_argument(
"--top_k", type=int, default=1, help="Top-k sampling parameter."
)
parser.add_argument(
"--output_image_path",
default="results",
help="Directory to save the annotated images.",
)
parser.add_argument(
"--input_json",
default="/Users/baixuehai/Downloads/2025_2/AITM_Test_General_BBox_v0.json",
help="Directory to save the annotated images.",
)
parser.add_argument(
"--output_json",
default="/Users/baixuehai/Downloads/2025_2/AITM_Test_General_BBox_v0.json",
help="Directory to save the annotated images.",
)
parser.add_argument(
"--format_key",
default="action_op_sensitive",
help="Key to select the prompt format.",
)
args = parser.parse_args()
# Dictionary mapping format keys to format strings
format_dict = {
"action_op_sensitive": "(Answer in Action-Operation-Sensitive format.)",
"status_plan_action_op": "(Answer in Status-Plan-Action-Operation format.)",
"status_action_op_sensitive": "(Answer in Status-Action-Operation-Sensitive format.)",
"status_action_op": "(Answer in Status-Action-Operation format.)",
"action_op": "(Answer in Action-Operation format.)",
}
# Ensure the provided format_key is valid
if args.format_key not in format_dict:
raise ValueError(
f"Invalid format_key. Available keys are: {list(format_dict.keys())}"
)
# Ensure the output directory exists
os.makedirs(args.output_image_path, exist_ok=True)
# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
args.model_dir,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
device_map="auto",
# quantization_config=BitsAndBytesConfig(load_in_8bit=True), # For INT8 quantization
# quantization_config=BitsAndBytesConfig(load_in_4bit=True), # For INT4 quantization
).eval()
# Initialize platform and selected format strings
platform_str = f"(Platform: {args.platform})\n"
format_str = format_dict[args.format_key]
# Initialize history lists
history_step = []
history_action = []
round_num = 1
# with open(args.input_json, "r") as f:
# data = json.load(f)
dataset = AITM_Dataset(args.input_json)
data_loader = DataLoader(dataset, batch_size=16, shuffle=False)
res = []
for x in tqdm(data_loader,desc="Processing items"):
#x = data[i]
img_path,task = x
image = []
for path in img_path:
image.append(Image.open(path).convert("RGB"))
#image = Image.open(img_path).convert("RGB")
#task = x['conversations'][0]['value']
# Verify history lengths match
if len(history_step) != len(history_action):
raise ValueError("Mismatch in lengths of history_step and history_action.")
# Format history steps for output
history_str = "\nHistory steps: "
for index, (step, action) in enumerate(zip(history_step, history_action)):
history_str += f"\n{index}. {step}\t{action}"
# Compose the query with task, platform, and selected format instructions
query = []
for x in task:
query.append(f"Task: {x}{history_str}\n{platform_str}{format_str}")
#query = f"Task: {task}{history_str}\n{platform_str}{format_str}"
#print(f"Round {round_num} query:\n{query}")
inputs = tokenizer.apply_chat_template(
[{"role": "user", "image": **image, "content": **query}],
add_generation_prompt=True,
tokenize=True,
return_tensors="pt",
return_dict=True,
).to(model.device)
# Generation parameters
gen_kwargs = {
"max_length": args.max_length,
"do_sample": True,
"top_k": args.top_k,
}
# Generate response
with torch.no_grad():
outputs = model.generate(**inputs, **gen_kwargs)
outputs = outputs[:, inputs["input_ids"].shape[1]:]
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
#print(f"Model response:\n{response}")
# Extract grounded operation and action
grounded_pattern = r"Grounded Operation:\s*(.*)"
action_pattern = r"Action:\s*(.*)"
matches_history = re.search(grounded_pattern, response)
matches_actions = re.search(action_pattern, response)
if matches_history:
grounded_operation = matches_history.group(1)
history_step.append(grounded_operation)
if matches_actions:
action_operation = matches_actions.group(1)
history_action.append(action_operation)
# Extract bounding boxes from the response
box_pattern = r"box=\[\[?(\d+),(\d+),(\d+),(\d+)\]?\]"
matches = re.findall(box_pattern, response)
if matches:
boxes = [[int(x) / 1000 for x in match] for match in matches]
# Extract base name of the user's input image (without extension)
base_name = []
for path in args.img_path:
base_name.append(os.path.splitext(os.path.basename(path))[0])
#base_name = os.path.splitext(os.path.basename(img_path))[0]
# Construct the output file name with round number
output_file_name = []
for i in range(len(base_name)):
output_file_name.append(f"{base_name[i]}_{round_num}_{i}.png")
#output_file_name = f"{base_name}_{round_num}.png"
output_path = []
for x in output_file_name:
output_path.append(os.path.join(args.output_image_path, x))
#output_path = os.path.join(args.output_image_path, output_file_name)
draw_boxes_on_image(image, boxes, output_path)
#print(f"Annotated image saved at: {output_path}")
ans = {
'query': f"Round {round_num} query:\n{query}",
'response': response,
'output_path': output_path
}
res.append(ans)
round_num += 1
#print(res)
print("Writing to json file")
with open(args.output_json, "w") as file:
print("Writing to json file")
json.dump(res, file, indent=4)
print("Done")
if __name__ == "__main__":
main() |