Albumentation库 bbox使用案例代码

import os
import numpy as np
import cv2
from matplotlib import pyplot as plt
from urllib.request import urlopen

from albumentations import (
    HorizontalFlip,
    VerticalFlip,
    Resize,
    CenterCrop,
    RandomCrop,
    Crop,
    Compose
)

# 用于图片上的边界框和类别 labels 的可视化函数
BOX_COLOR = (255, 0, 0)
TEXT_COLOR = (255, 255, 255)


def visualize_bbox(img, bbox, class_id, class_idx_to_name, color=BOX_COLOR, thickness=2):
    x_min, y_min, w, h = bbox
    x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color=color, thickness=thickness)
    class_name = class_idx_to_name[class_id]
    ((text_width, text_height), _) = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.35, 1)
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)), (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(img, class_name, (x_min, y_min - int(0.3 * text_height)), cv2.FONT_HERSHEY_SIMPLEX, 0.35, TEXT_COLOR,
                lineType=cv2.LINE_AA)
    return img


def get_aug(aug, min_area=0., min_visibility=0.):
    return Compose(aug, bbox_params={'format': 'coco', 'min_area': min_area,
                                     'min_visibility': min_visibility,
                                     'label_fields': ['category_id']})

image = cv2.imread('/home/yxq/桌面/mouth/data/images/5.png')

# Annotations for image 386298 from COCO http://cocodataset.org/#explore?id=386298
annotations = {'image': image,
               'bboxes': [[466., 179., 20., 20.]],
               'category_id': [1]}
category_id_to_name = {1: 'teeth',}

aug = get_aug([VerticalFlip(p=1)])
augmented = aug(**annotations)

bbox, img, cls_id = augmented['bboxes'], augmented['image'], augmented['category_id']

for index, box in enumerate(bbox):
    show_img = visualize_bbox(img, box, cls_id[index], category_id_to_name)

cv2.imshow('show_img', show_img)
cv2.waitKey(0)