|
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. |