Spaces:
Runtime error
Runtime error
Create helper.py
Browse files
helper.py
ADDED
@@ -0,0 +1,232 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import io
|
2 |
+
import matplotlib.pyplot as plt
|
3 |
+
import requests
|
4 |
+
import inflect
|
5 |
+
from PIL import Image
|
6 |
+
import torch
|
7 |
+
import numpy as np
|
8 |
+
|
9 |
+
def load_image_from_url(url):
|
10 |
+
return Image.open(requests.get(url, stream=True).raw)
|
11 |
+
|
12 |
+
def render_results_in_image(in_pil_img, in_results):
|
13 |
+
plt.figure(figsize=(16, 10))
|
14 |
+
plt.imshow(in_pil_img)
|
15 |
+
|
16 |
+
ax = plt.gca()
|
17 |
+
|
18 |
+
for prediction in in_results:
|
19 |
+
|
20 |
+
x, y = prediction['box']['xmin'], prediction['box']['ymin']
|
21 |
+
w = prediction['box']['xmax'] - prediction['box']['xmin']
|
22 |
+
h = prediction['box']['ymax'] - prediction['box']['ymin']
|
23 |
+
|
24 |
+
ax.add_patch(plt.Rectangle((x, y),
|
25 |
+
w,
|
26 |
+
h,
|
27 |
+
fill=False,
|
28 |
+
color="green",
|
29 |
+
linewidth=2))
|
30 |
+
ax.text(
|
31 |
+
x,
|
32 |
+
y,
|
33 |
+
f"{prediction['label']}: {round(prediction['score']*100, 1)}%",
|
34 |
+
color='red'
|
35 |
+
)
|
36 |
+
|
37 |
+
plt.axis("off")
|
38 |
+
|
39 |
+
# Save the modified image to a BytesIO object
|
40 |
+
img_buf = io.BytesIO()
|
41 |
+
plt.savefig(img_buf, format='png',
|
42 |
+
bbox_inches='tight',
|
43 |
+
pad_inches=0)
|
44 |
+
img_buf.seek(0)
|
45 |
+
modified_image = Image.open(img_buf)
|
46 |
+
|
47 |
+
# Close the plot to prevent it from being displayed
|
48 |
+
plt.close()
|
49 |
+
|
50 |
+
return modified_image
|
51 |
+
|
52 |
+
def summarize_predictions_natural_language(predictions):
|
53 |
+
summary = {}
|
54 |
+
p = inflect.engine()
|
55 |
+
|
56 |
+
for prediction in predictions:
|
57 |
+
label = prediction['label']
|
58 |
+
if label in summary:
|
59 |
+
summary[label] += 1
|
60 |
+
else:
|
61 |
+
summary[label] = 1
|
62 |
+
|
63 |
+
result_string = "In this image, there are "
|
64 |
+
for i, (label, count) in enumerate(summary.items()):
|
65 |
+
count_string = p.number_to_words(count)
|
66 |
+
result_string += f"{count_string} {label}"
|
67 |
+
if count > 1:
|
68 |
+
result_string += "s"
|
69 |
+
|
70 |
+
result_string += " "
|
71 |
+
|
72 |
+
if i == len(summary) - 2:
|
73 |
+
result_string += "and "
|
74 |
+
|
75 |
+
# Remove the trailing comma and space
|
76 |
+
result_string = result_string.rstrip(', ') + "."
|
77 |
+
|
78 |
+
return result_string
|
79 |
+
|
80 |
+
|
81 |
+
##### To ignore warnings #####
|
82 |
+
import warnings
|
83 |
+
import logging
|
84 |
+
from transformers import logging as hf_logging
|
85 |
+
|
86 |
+
def ignore_warnings():
|
87 |
+
# Ignore specific Python warnings
|
88 |
+
warnings.filterwarnings("ignore", message="Some weights of the model checkpoint")
|
89 |
+
warnings.filterwarnings("ignore", message="Could not find image processor class")
|
90 |
+
warnings.filterwarnings("ignore", message="The `max_size` parameter is deprecated")
|
91 |
+
|
92 |
+
# Adjust logging for libraries using the logging module
|
93 |
+
logging.basicConfig(level=logging.ERROR)
|
94 |
+
hf_logging.set_verbosity_error()
|
95 |
+
|
96 |
+
########
|
97 |
+
|
98 |
+
def show_mask(mask, ax, random_color=False):
|
99 |
+
if random_color:
|
100 |
+
color = np.concatenate([np.random.random(3),
|
101 |
+
np.array([0.6])],
|
102 |
+
axis=0)
|
103 |
+
else:
|
104 |
+
color = np.array([30/255, 144/255, 255/255, 0.6])
|
105 |
+
h, w = mask.shape[-2:]
|
106 |
+
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
107 |
+
ax.imshow(mask_image)
|
108 |
+
|
109 |
+
|
110 |
+
def show_box(box, ax):
|
111 |
+
x0, y0 = box[0], box[1]
|
112 |
+
w, h = box[2] - box[0], box[3] - box[1]
|
113 |
+
ax.add_patch(plt.Rectangle((x0, y0),
|
114 |
+
w,
|
115 |
+
h, edgecolor='green',
|
116 |
+
facecolor=(0,0,0,0),
|
117 |
+
lw=2))
|
118 |
+
|
119 |
+
def show_boxes_on_image(raw_image, boxes):
|
120 |
+
plt.figure(figsize=(10,10))
|
121 |
+
plt.imshow(raw_image)
|
122 |
+
for box in boxes:
|
123 |
+
show_box(box, plt.gca())
|
124 |
+
plt.axis('on')
|
125 |
+
plt.show()
|
126 |
+
|
127 |
+
def show_points_on_image(raw_image, input_points, input_labels=None):
|
128 |
+
plt.figure(figsize=(10,10))
|
129 |
+
plt.imshow(raw_image)
|
130 |
+
input_points = np.array(input_points)
|
131 |
+
if input_labels is None:
|
132 |
+
labels = np.ones_like(input_points[:, 0])
|
133 |
+
else:
|
134 |
+
labels = np.array(input_labels)
|
135 |
+
show_points(input_points, labels, plt.gca())
|
136 |
+
plt.axis('on')
|
137 |
+
plt.show()
|
138 |
+
|
139 |
+
def show_points_and_boxes_on_image(raw_image,
|
140 |
+
boxes,
|
141 |
+
input_points,
|
142 |
+
input_labels=None):
|
143 |
+
plt.figure(figsize=(10,10))
|
144 |
+
plt.imshow(raw_image)
|
145 |
+
input_points = np.array(input_points)
|
146 |
+
if input_labels is None:
|
147 |
+
labels = np.ones_like(input_points[:, 0])
|
148 |
+
else:
|
149 |
+
labels = np.array(input_labels)
|
150 |
+
show_points(input_points, labels, plt.gca())
|
151 |
+
for box in boxes:
|
152 |
+
show_box(box, plt.gca())
|
153 |
+
plt.axis('on')
|
154 |
+
plt.show()
|
155 |
+
|
156 |
+
|
157 |
+
def show_points_and_boxes_on_image(raw_image,
|
158 |
+
boxes,
|
159 |
+
input_points,
|
160 |
+
input_labels=None):
|
161 |
+
plt.figure(figsize=(10,10))
|
162 |
+
plt.imshow(raw_image)
|
163 |
+
input_points = np.array(input_points)
|
164 |
+
if input_labels is None:
|
165 |
+
labels = np.ones_like(input_points[:, 0])
|
166 |
+
else:
|
167 |
+
labels = np.array(input_labels)
|
168 |
+
show_points(input_points, labels, plt.gca())
|
169 |
+
for box in boxes:
|
170 |
+
show_box(box, plt.gca())
|
171 |
+
plt.axis('on')
|
172 |
+
plt.show()
|
173 |
+
|
174 |
+
|
175 |
+
def show_points(coords, labels, ax, marker_size=375):
|
176 |
+
pos_points = coords[labels==1]
|
177 |
+
neg_points = coords[labels==0]
|
178 |
+
ax.scatter(pos_points[:, 0],
|
179 |
+
pos_points[:, 1],
|
180 |
+
color='green',
|
181 |
+
marker='*',
|
182 |
+
s=marker_size,
|
183 |
+
edgecolor='white',
|
184 |
+
linewidth=1.25)
|
185 |
+
ax.scatter(neg_points[:, 0],
|
186 |
+
neg_points[:, 1],
|
187 |
+
color='red',
|
188 |
+
marker='*',
|
189 |
+
s=marker_size,
|
190 |
+
edgecolor='white',
|
191 |
+
linewidth=1.25)
|
192 |
+
|
193 |
+
|
194 |
+
def fig2img(fig):
|
195 |
+
"""Convert a Matplotlib figure to a PIL Image and return it"""
|
196 |
+
import io
|
197 |
+
buf = io.BytesIO()
|
198 |
+
fig.savefig(buf)
|
199 |
+
buf.seek(0)
|
200 |
+
img = Image.open(buf)
|
201 |
+
return img
|
202 |
+
|
203 |
+
|
204 |
+
def show_mask_on_image(raw_image, mask, return_image=False):
|
205 |
+
if not isinstance(mask, torch.Tensor):
|
206 |
+
mask = torch.Tensor(mask)
|
207 |
+
|
208 |
+
if len(mask.shape) == 4:
|
209 |
+
mask = mask.squeeze()
|
210 |
+
|
211 |
+
fig, axes = plt.subplots(1, 1, figsize=(15, 15))
|
212 |
+
|
213 |
+
mask = mask.cpu().detach()
|
214 |
+
axes.imshow(np.array(raw_image))
|
215 |
+
show_mask(mask, axes)
|
216 |
+
axes.axis("off")
|
217 |
+
plt.show()
|
218 |
+
|
219 |
+
if return_image:
|
220 |
+
fig = plt.gcf()
|
221 |
+
return fig2img(fig)
|
222 |
+
|
223 |
+
|
224 |
+
|
225 |
+
|
226 |
+
def show_pipe_masks_on_image(raw_image, outputs):
|
227 |
+
plt.imshow(np.array(raw_image))
|
228 |
+
ax = plt.gca()
|
229 |
+
for mask in outputs["masks"]:
|
230 |
+
show_mask(mask, ax=ax, random_color=True)
|
231 |
+
plt.axis("off")
|
232 |
+
plt.show()
|