utils/augmentation.py

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)