File size: 4,510 Bytes
c426a27
 
2461d7d
 
c426a27
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2461d7d
c426a27
 
 
 
 
 
 
 
 
10240e0
c426a27
2461d7d
c426a27
 
 
 
 
 
5c74464
 
c426a27
 
 
 
 
 
 
 
2461d7d
c426a27
 
 
5c74464
2461d7d
 
c426a27
 
 
2461d7d
c426a27
2461d7d
 
 
 
 
 
 
 
 
 
 
 
 
 
c426a27
 
 
 
 
 
 
 
2461d7d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from PIL import Image, ImageDraw, ImageFont
import copy
import numpy as np
import cv2 

def wrap_text(text, font, max_width):
    lines = []
    words = text.split(' ')
    current_line = ''

    for word in words:
        if font.getsize(current_line + word)[0] <= max_width:
            current_line += word + ' '
        else:
            lines.append(current_line)
            current_line = word + ' '

    lines.append(current_line)
    return lines

def create_bubble_frame(image, text, point, segmask, input_points, input_labels, font_path='times_with_simsun.ttf', font_size_ratio=0.033, point_size_ratio=0.01):
    # Load the image
    if type(image) == np.ndarray:
        image = Image.fromarray(image)
        
    image = copy.deepcopy(image)
    width, height = image.size

    # Calculate max_text_width and font_size based on image dimensions and total number of characters
    total_chars = len(text)
    max_text_width = int(0.4 * width)
    font_size = int(height * font_size_ratio)
    point_size = max(int(height * point_size_ratio), 1)

    # Load the font
    font = ImageFont.truetype(font_path, font_size)

    # Wrap the text to fit within the max_text_width
    lines = wrap_text(text, font, max_text_width)
    text_width = max([font.getsize(line)[0] for line in lines])
    _, text_height = font.getsize(lines[0])
    text_height = text_height * len(lines)

    # Define bubble frame dimensions
    padding = 10
    bubble_width = text_width + 2 * padding
    bubble_height = text_height + 2 * padding

    # Create a new image for the bubble frame
    bubble = Image.new('RGBA', (bubble_width, bubble_height), (255,248, 220, 0))

    # Draw the bubble frame on the new image
    draw = ImageDraw.Draw(bubble)
    # draw.rectangle([(0, 0), (bubble_width - 1, bubble_height - 1)], fill=(255, 255, 255, 0), outline=(255, 255, 255, 0), width=2)
    draw_rounded_rectangle(draw, (0, 0, bubble_width - 1, bubble_height - 1), point_size * 2, 
                        fill=(255,248, 220, 120), outline=None, width=2)
    # Draw the wrapped text line by line
    y_text = padding
    for line in lines:
        draw.text((padding, y_text), line, font=font, fill=(0, 0, 0, 255))
        y_text += font.getsize(line)[1]
        
    # Determine the point by the min area rect of mask
    try:
        ret, thresh = cv2.threshold(segmask, 127, 255, 0)
        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        largest_contour = max(contours, key=cv2.contourArea)
        min_area_rect = cv2.minAreaRect(largest_contour)
        box = cv2.boxPoints(min_area_rect)
        sorted_points = box[np.argsort(box[:, 0])]
        right_most_points = sorted_points[-2:]
        right_down_most_point = right_most_points[np.argsort(right_most_points[:, 1])][-1]
        x, y = int(right_down_most_point[0]), int(right_down_most_point[1])
    except:
        x, y = point
    # Calculate the bubble frame position
    if x + bubble_width > width:
        x = width - bubble_width
    if y + bubble_height > height:
        y = height - bubble_height

    # Paste the bubble frame onto the image
    image.paste(bubble, (x, y), bubble)
    draw = ImageDraw.Draw(image)
    colors = [(0, 191, 255, 255), (255, 106, 106, 255)]
    for p, label in zip(input_points, input_labels):
        point_x, point_y = p[0], p[1]
        left = point_x - point_size
        top = point_y - point_size
        right = point_x + point_size
        bottom = point_y + point_size
        draw.ellipse((left, top, right, bottom), fill=colors[label])
    return image


def draw_rounded_rectangle(draw, xy, corner_radius, fill=None, outline=None, width=1):
    x1, y1, x2, y2 = xy

    draw.rectangle(
        (x1, y1 + corner_radius, x2, y2 - corner_radius),
        fill=fill,
        outline=outline,
        width=width
    )
    draw.rectangle(
        (x1 + corner_radius, y1, x2 - corner_radius, y2),
        fill=fill,
        outline=outline,
        width=width
    )

    draw.pieslice((x1, y1, x1 + corner_radius * 2, y1 + corner_radius * 2), 180, 270, fill=fill, outline=outline, width=width)
    draw.pieslice((x2 - corner_radius * 2, y1, x2, y1 + corner_radius * 2), 270, 360, fill=fill, outline=outline, width=width)
    draw.pieslice((x2 - corner_radius * 2, y2 - corner_radius * 2, x2, y2), 0, 90, fill=fill, outline=outline, width=width)
    draw.pieslice((x1, y2 - corner_radius * 2, x1 + corner_radius * 2, y2), 90, 180, fill=fill, outline=outline, width=width)