Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4from PIL import Image, ImageDraw, ImageFont
5import re
6
7
8def draw_circle(draw, center, radius=10, width=2, outline_color=(0, 255, 0), is_fill=False, bg_color=(0, 255, 0), transparency=80):
9 # Calculate the bounding box coordinates for the circle
10 x1 = center[0] - radius
11 y1 = center[1] - radius
12 x2 = center[0] + radius
13 y2 = center[1] + radius
14 bbox = (x1, y1, x2, y2)
15
16 # Draw the circle
17 if is_fill:
18 # Calculate the alpha value based on the transparency percentage
19 alpha = int((1 - transparency / 100) * 255)
20
21 # Set the fill color with the specified background color and transparency
22 fill_color = tuple(bg_color) + (alpha,)
23
24 draw.ellipse(bbox, width=width, outline=outline_color, fill=fill_color)
25 else:
26 draw.ellipse(bbox, width=width, outline=outline_color)
27
28def draw_point(draw, center, radius1=3, radius2=6, color=(0, 255, 0)):
29 draw_circle(draw, center, radius=radius1, outline_color=color)
30 draw_circle(draw, center, radius=radius2, outline_color=color)
31
32def draw_rectangle(draw, box_coords, width=2, outline_color=(0, 255, 0), is_fill=False, bg_color=(0, 255, 0), transparency=80):
33 if is_fill:
34 # Calculate the alpha value based on the transparency percentage
35 alpha = int((1 - transparency / 100) * 255)
36
37 # Set the fill color with the specified background color and transparency
38 fill_color = tuple(bg_color) + (alpha,)
39
40 draw.rectangle(box_coords, width=width, outline=outline_color, fill=fill_color)
41 else:
42 draw.rectangle(box_coords, width=width, outline=outline_color)
43
44def draw(path, out_path, response):
45 img = Image.open(path).convert("RGB")
46 draw = ImageDraw.Draw(img)
47
48 box_coords = re.findall(r"<box>(.*?)</box>", response)
49 for box in box_coords:
50 try:
51 x1, y1, x2, y2 = box.replace("(", "").replace(")", "").split(",")
52 x1, y1, x2, y2 = float(x1) * img.width/1000, float(y1) * img.height/1000, float(x2) * img.width/1000, float(y2) * img.height/1000
53 draw_rectangle(draw, (x1, y1, x2, y2))
54 except:
55 print("There were some errors while parsing the bounding box.")
56
57 point_coords = re.findall(r"<point>(.*?)</point>", response)
58 for point in point_coords:
59 try:
60 x1, y1 = point.replace("(", "").replace(")", "").split(",")
61 x1, y1 = float(x1) * img.width/1000, float(y1) * img.height/1000
62 draw_point(draw, (x1, y1))
63 except:
64 print("There were some errors while parsing the bounding point.")
65
66 img.save(out_path)
67
68def load_model_and_tokenizer(path, device):
69 tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
70 model = AutoModelForCausalLM.from_pretrained(path, device_map=device, trust_remote_code=True).eval()
71 return model, tokenizer
72
73
74def infer(model, tokenizer, image_path, text):
75 query = tokenizer.from_list_format([
76 {'image': image_path},
77 {'text': text},
78 ])
79 response, history = model.chat(tokenizer, query=query, history=None)
80 return response
81
82if __name__ == "__main__":
83 device = "cuda:0"
84 model_path = "<your_model_path>"
85 model, tokenizer = load_model_and_tokenizer(model_path, device)
86
87 while True:
88 image_path = input("image path >>>>> ")
89 if image_path == "stop":
90 break
91 query = input("Human:")
92 if query == "stop":
93 break
94
95 response = infer(model, tokenizer, image_path, query)
96 draw(image_path, "1.jpg", response)