MSaadTariq commited on
Commit
8371896
·
verified ·
1 Parent(s): c22784b

Create helper.py

Browse files
Files changed (1) hide show
  1. helper.py +232 -0
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()