File size: 1,186 Bytes
5fa1a76 |
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 |
import numpy as np
import os
from PIL import Image, ImageDraw
image = cppe5["train"][0]["image"]
annotations = cppe5["train"][0]["objects"]
draw = ImageDraw.Draw(image)
categories = cppe5["train"].features["objects"].feature["category"].names
id2label = {index: x for index, x in enumerate(categories, start=0)}
label2id = {v: k for k, v in id2label.items()}
for i in range(len(annotations["id"])):
box = annotations["bbox"][i]
class_idx = annotations["category"][i]
x, y, w, h = tuple(box)
# Check if coordinates are normalized or not
if max(box) > 1.0:
# Coordinates are un-normalized, no need to re-scale them
x1, y1 = int(x), int(y)
x2, y2 = int(x + w), int(y + h)
else:
# Coordinates are normalized, re-scale them
x1 = int(x * width)
y1 = int(y * height)
x2 = int((x + w) * width)
y2 = int((y + h) * height)
draw.rectangle((x, y, x + w, y + h), outline="red", width=1)
draw.text((x, y), id2label[class_idx], fill="white")
image
To visualize the bounding boxes with associated labels, you can get the labels from the dataset's metadata, specifically
the category field. |