42 lines
1.8 KiB
Python
42 lines
1.8 KiB
Python
|
import torch
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
import numpy as np
|
|||
|
|
import albumentations as A
|
|||
|
|
|
|||
|
|
# boxes = (cls, x, y, w, h)
|
|||
|
|
def horizontal_flip(images, boxes):
|
|||
|
|
images = np.flip(images, [-1])
|
|||
|
|
boxes[:, 1] = 1 - boxes[:, 1]
|
|||
|
|
return images, boxes
|
|||
|
|
|
|||
|
|
# images[np.unit8], boxes[numpy] = (cls, x, y, w, h)
|
|||
|
|
def augment(image, boxes):
|
|||
|
|
h, w, _ = image.shape
|
|||
|
|
labels, boxes_coord = boxes[:, 0], boxes[:, 1:]
|
|||
|
|
labels = labels.tolist()
|
|||
|
|
boxes_coord = boxes_coord * h # 得到原图尺寸下的坐标（未归一化的坐标）
|
|||
|
|
boxes_coord[:, 0] = np.clip(boxes_coord[:, 0]-boxes_coord[:, 2]/2, a_min=0, a_max=None) # 确保x_min和y_min有效
|
|||
|
|
boxes_coord[:, 1] = np.clip(boxes_coord[:, 1]-boxes_coord[:, 3]/2, a_min=0, a_max=None)
|
|||
|
|
boxes_coord = boxes_coord.tolist() # [x_min, y_min, width, height]
|
|||
|
|
|
|||
|
|
# 在这里设置数据增强的方法
|
|||
|
|
aug = A.Compose([
|
|||
|
|
A.HorizontalFlip(p=0.5),
|
|||
|
|
# A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, p=0.5),
|
|||
|
|
# A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=5, border_mode=0, p=0.5)
|
|||
|
|
], bbox_params={'format':'coco', 'label_fields': ['category_id']})
|
|||
|
|
|
|||
|
|
augmented = aug(image=image, bboxes=boxes_coord, category_id=labels)
|
|||
|
|
|
|||
|
|
# 经过aug之后，如果把boxes变没了，则返回原来的图片
|
|||
|
|
if augmented['bboxes']:
|
|||
|
|
image = augmented['image']
|
|||
|
|
|
|||
|
|
boxes_coord = np.array(augmented['bboxes']) # x_min, y_min, w, h → x, y, w, h
|
|||
|
|
boxes_coord[:, 0] = boxes_coord[:, 0] + boxes_coord[:, 2]/2
|
|||
|
|
boxes_coord[:, 1] = boxes_coord[:, 1] + boxes_coord[:, 3]/2
|
|||
|
|
boxes_coord = boxes_coord / h
|
|||
|
|
labels = np.array(augmented['category_id'])[:, None]
|
|||
|
|
boxes = np.concatenate((labels, boxes_coord), 1)
|
|||
|
|
|
|||
|
|
return image, boxes
|