File size: 8,135 Bytes
320c5b2 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import cv2
from cv2 import resize
import numpy as np
import random
class DataAugmentor(object):
def __init__(self, cfg, mode="train"):
self.cfg = cfg
self.mean = np.asarray([0.48145466, 0.4578275,
0.40821073], dtype=np.float32)
self.std = np.asarray([0.26862954, 0.26130258,
0.27577711], dtype=np.float32)
self.mode = mode
def _random_mirror(self, data_dict):
img = data_dict["rgb"]
bboxes = data_dict["bboxes"][:, :4]
if random.randint(0, 1):
_, width, _ = img.shape
data_dict["rgb"] = img[:, ::-1]
data_dict["depth"] = data_dict["depth"][:, ::-1]
data_dict["ins_masks"] = data_dict["ins_masks"][:, :, ::-1]
data_dict["grasp_masks"]["qua"] = data_dict["grasp_masks"]["qua"][:, :, ::-1]
data_dict["grasp_masks"]["ang"] = data_dict["grasp_masks"]["ang"][:, :, ::-1]
data_dict["grasp_masks"]["wid"] = data_dict["grasp_masks"]["wid"][:, :, ::-1]
bboxes[:, 0::2] = width - bboxes[:, 2::-2]
data_dict["bboxes"][:, :4] = bboxes
def _random_brightness(self, img, delta=32):
img += random.uniform(-delta, delta)
return np.clip(img, 0., 255.)
def _random_contrast(self, img, lower=0.7, upper=1.3):
img *= random.uniform(lower, upper)
return np.clip(img, 0., 255.)
def _random_saturation(self, img, delta=15.):
img[:, :, 0] += random.uniform(-delta, delta)
img[:, :, 0][img[:, :, 0] > 360.0] -= 360.0
img[:, :, 0][img[:, :, 0] < 0.0] += 360.0
return img
def _random_hue(self, img, delta=15.):
img[:, :, 0] += random.uniform(-delta, delta)
img[:, :, 0][img[:, :, 0] > 360.0] -= 360.0
img[:, :, 0][img[:, :, 0] < 0.0] += 360.0
return img
def _photometric_distort(self, data_dict):
img = data_dict["rgb"].astype(np.float32)
if random.randint(0, 1):
img = self._random_brightness(img)
if random.randint(0, 1):
img = self._random_contrast(img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
img = self._random_saturation(img)
img = self._random_hue(img)
img = cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
img = np.clip(img, 0., 255.)
data_dict["rgb"] = img
def _pad_to_square(self, data_dict):
img = data_dict["rgb"]
img_h, img_w = img.shape[:2]
if img_h == img_w:
pass
else:
pad_size = max(img_h, img_w)
pad_img = np.zeros((pad_size, pad_size, 3), dtype='float32')
pad_img[:, :, :] = self.mean
pad_depth = np.zeros((pad_size, pad_size), dtype='float32')
pad_ins_masks = np.zeros((data_dict["ins_masks"].shape[0], pad_size, pad_size), dtype='float32')
pad_qua_masks = np.zeros((data_dict["grasp_masks"]["qua"].shape[0], pad_size, pad_size), dtype='float32')
pad_ang_masks = np.zeros((data_dict["grasp_masks"]["ang"].shape[0], pad_size, pad_size), dtype='float32')
pad_wid_masks = np.zeros((data_dict["grasp_masks"]["wid"].shape[0], pad_size, pad_size), dtype='float32')
if self.mode=="train":
if img_h < img_w:
random_y1 = random.randint(0, img_w - img_h)
pad_img[random_y1: random_y1 + img_h, :, :] = data_dict["rgb"]
pad_depth[random_y1: random_y1 + img_h, :] = data_dict["depth"]
pad_ins_masks[:, random_y1: random_y1 + img_h, :] = data_dict["ins_masks"]
pad_qua_masks[:, random_y1: random_y1 + img_h, :] = data_dict["grasp_masks"]["qua"]
pad_ang_masks[:, random_y1: random_y1 + img_h, :] = data_dict["grasp_masks"]["ang"]
pad_wid_masks[:, random_y1: random_y1 + img_h, :] = data_dict["grasp_masks"]["wid"]
data_dict["bboxes"][:, [1, 3]] += random_y1
if img_h > img_w:
random_x1 = random.randint(0, img_h - img_w)
pad_img[:, random_x1: random_x1 + img_w, :] = data_dict["rgb"]
pad_depth[:, random_x1: random_x1 + img_w] = data_dict["depth"]
pad_ins_masks[:, :, random_x1: random_x1 + img_w] = data_dict["ins_masks"]
pad_qua_masks[:, :, random_x1: random_x1 + img_w] = data_dict["grasp_masks"]["qua"]
pad_ang_masks[:, :, random_x1: random_x1 + img_w] = data_dict["grasp_masks"]["ang"]
pad_wid_masks[:, :, random_x1: random_x1 + img_w] = data_dict["grasp_masks"]["wid"]
data_dict["bboxes"][:, [0, 2]] += random_x1
elif self.mode in ["test", "val"]:
pad_img[0: img_h, 0: img_w, :] = data_dict["rgb"]
pad_depth[0: img_h, 0: img_w] = data_dict["depth"]
pad_ins_masks[:, 0: img_h, 0: img_w] = data_dict["ins_masks"]
pad_qua_masks[:, 0: img_h, 0: img_w] = data_dict["grasp_masks"]["qua"]
pad_ang_masks[:, 0: img_h, 0: img_w] = data_dict["grasp_masks"]["ang"]
pad_wid_masks[:, 0: img_h, 0: img_w] = data_dict["grasp_masks"]["wid"]
data_dict["rgb"] = pad_img
data_dict["depth"] = pad_depth
data_dict["ins_masks"] = pad_ins_masks
data_dict["grasp_masks"]["qua"] = pad_qua_masks
data_dict["grasp_masks"]["ang"] = pad_ang_masks
data_dict["grasp_masks"]["wid"] = pad_wid_masks
def _resize(self, data_dict):
ori_size = data_dict["rgb"].shape[0]
tgt_size = self.cfg.img_size
scale = tgt_size / ori_size
data_dict["rgb"] = cv2.resize(data_dict["rgb"], (tgt_size, tgt_size))
data_dict["depth"] = cv2.resize(data_dict["depth"], (tgt_size, tgt_size))
data_dict["ins_masks"] = cv2.resize(data_dict["ins_masks"].transpose((1,2,0)), (tgt_size, tgt_size)).transpose((2,0,1)) if data_dict["ins_masks"].shape[0] > 1 else \
np.expand_dims(cv2.resize(data_dict["ins_masks"].transpose((1,2,0)), (tgt_size, tgt_size)), 0).transpose((2,0,1))
data_dict["grasp_masks"]["qua"] = cv2.resize(data_dict["grasp_masks"]["qua"].transpose((1,2,0)), (tgt_size, tgt_size)).transpose((2,0,1)) if data_dict["grasp_masks"]["qua"].shape[0] > 1 else \
np.expand_dims(cv2.resize(data_dict["grasp_masks"]["qua"].transpose((1,2,0)), (tgt_size, tgt_size)), 0).transpose((2,0,1))
data_dict["grasp_masks"]["ang"] = cv2.resize(data_dict["grasp_masks"]["ang"].transpose((1,2,0)), (tgt_size, tgt_size)).transpose((2,0,1)) if data_dict["grasp_masks"]["ang"].shape[0] > 1 else \
np.expand_dims(cv2.resize(data_dict["grasp_masks"]["ang"].transpose((1,2,0)), (tgt_size, tgt_size)), 0).transpose((2,0,1))
data_dict["grasp_masks"]["wid"] = cv2.resize(data_dict["grasp_masks"]["wid"].transpose((1,2,0)), (tgt_size, tgt_size)).transpose((2,0,1)) if data_dict["grasp_masks"]["wid"].shape[0] > 1 else \
np.expand_dims(cv2.resize(data_dict["grasp_masks"]["wid"].transpose((1,2,0)), (tgt_size, tgt_size)), 0).transpose((2,0,1))
data_dict["bboxes"][:, :4] *= scale
def _normalize_boxes(self, data_dict):
h, w = data_dict["rgb"].shape[:2]
data_dict["bboxes"][:, [0, 2]] /= w
data_dict["bboxes"][:, [1, 3]] /= h
def _normalize_img(self, data_dict):
img = data_dict["rgb"] / 255.
# img = (img - self.mean) / self.std
img = img[:, :, (2,1,0)]
img = np.transpose(img, (2, 0, 1))
data_dict["rgb"] = img
def __call__(self, data_dict):
if self.mode == "train":
self._photometric_distort(data_dict)
self._random_mirror(data_dict)
self._pad_to_square(data_dict)
self._resize(data_dict)
self._normalize_boxes(data_dict)
self._normalize_img(data_dict) |