-
Notifications
You must be signed in to change notification settings - Fork 0
/
detect_contours.py
137 lines (105 loc) · 4.7 KB
/
detect_contours.py
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
import cv2
import numpy as np
from PIL import Image
def create_final_image(image, dialouge):
# print("in create final image")
image1 = create_image_with_text(image, dialouge)
# print("created simage with text")
if image1 is None:
fallback_image = np.array(Image.open("image3.png"))
# print("in fa/llback")
return create_image_with_text(fallback_image, dialouge)
# print("returning image")
return image1
def create_image_with_text(image, dialouge):
# Edge detection
edges = cv2.Canny(image, 100, 200)
# print(edges)
# Contour detection
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Filter contours by size and aspect ratio
thought_bubbles = []
for contour in contours:
area = cv2.contourArea(contour)
area = area / (image.shape[0] * image.shape[1])
if area > 1e-2 and area < 0.5 :
x, y, w, h = cv2.boundingRect(contour)
aspect_ratio = float(w) / h
if aspect_ratio > 1.2 :
thought_bubbles.append(contour)
# print(len(thought_bubbles))
# if len(thought_bubbles) != 4 :
# return None
if len(thought_bubbles) == 4:
center_coords = []
for i, contour in enumerate(thought_bubbles):
mask = np.zeros_like(image)
# Draw filled contours on the mask
cv2.drawContours(mask, [thought_bubbles[i]], -1, (255, 255, 255), thickness=cv2.FILLED)
# create a binary image of the mask
mask2 = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
# Calculate the center of mass (centroid) of the white regions in the mask
M = cv2.moments(mask2, binaryImage=True)
center_x = int(M['m10'] / M['m00'])
center_y = int(M['m01'] / M['m00'])
center_coords.append((center_x, center_y))
x_coords = [x for x, y in center_coords]
y_coords = [y for x, y in center_coords]
sorted_x_coords = np.argsort(x_coords)
i1 = sorted_x_coords[0]
i2 = sorted_x_coords[1]
y1 = y_coords[i1]
y2 = y_coords[i2]
if y1 < y2:
top_left = i1
bottom_left = i2
else:
top_left = i2
bottom_left = i1
i3 = sorted_x_coords[2]
i4 = sorted_x_coords[3]
if y_coords[i3] < y_coords[i4]:
top_right = i3
bottom_right = i4
else:
top_right = i4
bottom_right = i3
thought_bubbles = [thought_bubbles[top_left], thought_bubbles[top_right], thought_bubbles[bottom_left], thought_bubbles[bottom_right]]
else:
thought_bubbles = thought_bubbles[:4]
for i, contour in enumerate(thought_bubbles):
text = dialouge[i]
mask = np.zeros_like(image)
# Draw filled contours on the mask
cv2.drawContours(mask, [thought_bubbles[i]], -1, (255, 255, 255), thickness=cv2.FILLED)
# create a binary image of the mask
mask2 = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
# Calculate the center of mass (centroid) of the white regions in the mask
M = cv2.moments(mask2, binaryImage=True)
center_x = int(M['m10'] / M['m00'])
center_y = int(M['m01'] / M['m00'])
# Define multiple lines of text
# split the text after every 4 spaces
lines_of_text = text.split(' ')
lines_of_text = [' '.join(lines_of_text[i:i+4]) for i in range(0, len(lines_of_text), 4)]
# lines_of_text = ['']
# Calculate the font scale based on the mask size
font_scale = min(mask.shape[0], mask.shape[1]) / 2000 # Adjust the divisor as needed
font_color = (255, 0, 0) # Blue color in BGR
font_thickness = 2
font = cv2.FONT_HERSHEY_SIMPLEX
# Calculate the total height of all lines of text
total_text_height = len(lines_of_text) * cv2.getTextSize(lines_of_text[0], font, font_scale, font_thickness)[0][1]
# Calculate the starting y-coordinate to center the text vertically
start_y = center_y - total_text_height // 2
# Add each line of text to the mask
for line in lines_of_text:
# Calculate the size of the current line of text
text_size = cv2.getTextSize(line, font, font_scale, font_thickness)[0]
# Calculate the x-coordinate to center the text horizontally
text_x = center_x - text_size[0] // 2
# Add the text to the mask
cv2.putText(image, line, (text_x, start_y), font, font_scale, font_color, font_thickness)
# Move the starting y-coordinate down for the next line
start_y += text_size[1]
return image