Spaces:
Sleeping
Sleeping
File size: 2,677 Bytes
14f6839 07cb7d3 | 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 | 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)
|