| 152 | |
| 153 | |
| 154 | def visualize_detections(image_path, output_path, detections, labels=[]): |
| 155 | image = Image.open(image_path).convert(mode="RGB") |
| 156 | draw = ImageDraw.Draw(image) |
| 157 | line_width = 2 |
| 158 | font = ImageFont.load_default() |
| 159 | for d in detections: |
| 160 | color = COLORS[d["class"] % len(COLORS)] |
| 161 | draw.line( |
| 162 | [ |
| 163 | (d["xmin"], d["ymin"]), |
| 164 | (d["xmin"], d["ymax"]), |
| 165 | (d["xmax"], d["ymax"]), |
| 166 | (d["xmax"], d["ymin"]), |
| 167 | (d["xmin"], d["ymin"]), |
| 168 | ], |
| 169 | width=line_width, |
| 170 | fill=color, |
| 171 | ) |
| 172 | label = "Class {}".format(d["class"]) |
| 173 | if d["class"] < len(labels): |
| 174 | label = "{}".format(labels[d["class"]]) |
| 175 | score = d["score"] |
| 176 | text = "{}: {}%".format(label, int(100 * score)) |
| 177 | if score < 0: |
| 178 | text = label |
| 179 | left, top, right, bottom = font.getbbox(text) |
| 180 | text_width, text_height = right - left, bottom - top |
| 181 | text_bottom = max(text_height, d["ymin"]) |
| 182 | text_left = d["xmin"] |
| 183 | margin = np.ceil(0.05 * text_height) |
| 184 | draw.rectangle( |
| 185 | [(text_left, text_bottom - text_height - 2 * margin), (text_left + text_width, text_bottom)], fill=color |
| 186 | ) |
| 187 | draw.text((text_left + margin, text_bottom - text_height - margin), text, fill="black", font=font) |
| 188 | if output_path is None: |
| 189 | return image |
| 190 | image.save(output_path) |
| 191 | |
| 192 | |
| 193 | def concat_visualizations(images, names, colors, output_path): |