Spaces:
Sleeping
Sleeping
| from importlib import import_module | |
| from io import BytesIO | |
| import matplotlib.pyplot as plt | |
| from matplotlib.colors import hsv_to_rgb | |
| from matplotlib.patches import Rectangle | |
| from PIL import Image | |
| color_map = {0: "gray"} | |
| def label_to_color(label): | |
| if label not in color_map: | |
| h, s, v = (len(color_map) * 0.618) % 1, 0.5, 0.9 | |
| color_map[label] = hsv_to_rgb((h, s, v)) | |
| return color_map[label] | |
| def draw_box(ax, box, text, color): | |
| x, y, w, h = box | |
| ax.add_patch(Rectangle((x, y), w, h, lw=2, ec=color, fc="none")) | |
| textbox = dict(fc=color, pad=1, ec="none") | |
| ax.text(x, y, text, c="white", size=10, va="bottom", bbox=textbox) | |
| def draw_image(ax, image): | |
| ax.set(xlim=(0, 1), ylim=(1, 0), xticks=[], yticks=[], aspect="equal") | |
| height, width = image.shape[:2] | |
| hpad = (1 - height / width) / 2 if width > height else 0 | |
| wpad = (1 - width / height) / 2 if height > width else 0 | |
| extent = [wpad, 1 - wpad, 1 - hpad, hpad] | |
| ax.imshow(image, extent=extent) | |
| def from_grid(loc, box, grid_size: int): | |
| (xi, yi), (x, y, w, h) = loc, box | |
| x = (xi + x) / grid_size - w / 2 | |
| y = (yi + y) / grid_size - h / 2 | |
| return x, y, w, h | |
| def draw_prediction( | |
| image, | |
| boxes, | |
| classes, | |
| grid_size: int, | |
| cutoff=None | |
| ) -> Image.Image: | |
| fig, ax = plt.subplots(dpi=150) | |
| draw_image(ax, image) | |
| for yi, row in enumerate(classes): | |
| for xi, label in enumerate(row): | |
| color = label_to_color(label) if label else "none" | |
| x, y, w, h = (v / grid_size for v in (xi, yi, 1.0, 1.0)) | |
| rect = Rectangle((x, y), w, h, lw=2, ec="black", fc=color, alpha=0.5) | |
| ax.add_patch(rect) | |
| for yi, row in enumerate(boxes): | |
| for xi, box in enumerate(row): | |
| box, confidence = box[:4], box[4] | |
| if not cutoff or confidence >= cutoff: | |
| normalized_box = from_grid((xi, yi), box, grid_size) | |
| label = int(classes[yi, xi]) | |
| color = label_to_color(label) | |
| name = resolve_label_name(label) | |
| draw_box(ax, normalized_box, f"{name} {max(confidence, 0):.2f}", color) | |
| buffer = BytesIO() | |
| fig.savefig(buffer, format="png") | |
| plt.close(fig) | |
| buffer.seek(0) | |
| return Image.open(buffer).convert("RGB") | |
| def resolve_label_name(label: int) -> str: | |
| try: | |
| keras_hub = import_module("keras_hub") | |
| except ModuleNotFoundError: | |
| return f"label{label}" | |
| if not hasattr(keras_hub, "utils"): | |
| return f"label{label}" | |
| if not hasattr(keras_hub.utils, "coco_id_to_name"): | |
| return f"label{label}" | |
| return keras_hub.utils.coco_id_to_name(label) | |