RuiTerrty/RemoteSensingChangeDetection-RSCD.HA2F

1import numpy2import numpy as np3import torch4import random5import cv26 7 8class Scale(object):9    """10    Resize the given image to a fixed scale11    """12 13    def __init__(self, wi, he):14        '''15        :param wi: width after resizing16        :param he: height after reszing17        '''18        self.w = wi19        self.h = he20 21    # modified from torchvision to add support for max size22 23    def __call__(self, img, label):24        '''25        :param img: RGB image26        :param label: semantic label image27        :return: resized images28        '''29        # bilinear interpolation for RGB image30        img = cv2.resize(img, (self.w, self.h))31        # nearest neighbour interpolation for label image32        label = cv2.resize(label, (self.w, self.h), interpolation=cv2.INTER_NEAREST)33        return [img, label]34 35 36class Resize(object):37    def __init__(self, min_size, max_size, strict=False):38        if not isinstance(min_size, (list, tuple)):39            min_size = (min_size,)40        self.min_size = min_size41        self.max_size = max_size42        self.strict = strict43 44    # modified from torchvision to add support for max size45    def get_size(self, image_size):46        w, h = image_size47        if not self.strict:48            size = random.choice(self.min_size)49            max_size = self.max_size50            if max_size is not None:51                min_original_size = float(min((w, h)))52                max_original_size = float(max((w, h)))53                if max_original_size / min_original_size * size > max_size:54                    size = int(round(max_size * min_original_size / max_original_size))55 56            if (w <= h and w == size) or (h <= w and h == size):57                return (h, w)58 59            if w < h:60                ow = size61                oh = int(size * h / w)62            else:63                oh = size64                ow = int(size * w / h)65 66            return (oh, ow)67        else:68            if w < h:69                return (self.max_size, self.min_size[0])70            else:71                return (self.min_size[0], self.max_size)72 73    def __call__(self, image, label):74        size = self.get_size(image.shape[:2])75        image = cv2.resize(image, size)76        # I confirm that the output size is right, not reversed77        label = cv2.resize(label, size, interpolation=cv2.INTER_NEAREST)78        return (image, label)79 80 81class RandomCropResize(object):82    """83    Randomly crop and resize the given image with a probability of 0.584    """85 86    def __init__(self, crop_area):87        '''88        :param crop_area: area to be cropped (this is the max value and we select between 0 and crop area89        '''90        self.cw = crop_area91        self.ch = crop_area92 93    def __call__(self, img, label):94        if random.random() < 0.5:95            h, w = img.shape[:2]96            x1 = random.randint(0, self.ch)97            y1 = random.randint(0, self.cw)98 99            img_crop = img[y1:h - y1, x1:w - x1]100            label_crop = label[y1:h - y1, x1:w - x1]101 102            img_crop = cv2.resize(img_crop, (w, h))103            label_crop = cv2.resize(label_crop, (w, h), interpolation=cv2.INTER_NEAREST)104 105            return img_crop, label_crop106        else:107            return [img, label]108 109 110class RandomFlip(object):111    """112    Randomly flip the given Image with a probability of 0.5113    """114 115    def __call__(self, image, label):116        if random.random() < 0.5:117                image = cv2.flip(image, 0)  # horizontal flip118                label = cv2.flip(label, 0)  # horizontal flip119        if random.random() < 0.5:120                image = cv2.flip(image, 1)  # veritcal flip121                label = cv2.flip(label, 1)  # veritcal flip122        return [image, label]123 124 125class RandomExchange(object):126    """127    Randomly flip the given Image with a probability of 0.5128    """129 130    def __call__(self, image, label):131        if random.random() < 0.5:132            pre_img = image[:, :, 0:3]133            post_img = image[:, :, 3:6]134            image = numpy.concatenate((post_img, pre_img), axis=2)135        return [image, label]136 137 138class Normalize(object):139    """140    Given mean: (B, G, R) and std: (B, G, R),141    will normalize each channel of the torch.*Tensor, i.e.142    channel = (channel - mean) / std143    """144 145    def __init__(self, mean, std):146        '''147        :param mean: global mean computed from dataset148        :param std: global std computed from dataset149        '''150        self.mean = mean151        self.std = std152        self.depth_mean = [0.5]153        self.depth_std = [0.5]154 155    def __call__(self, image, label):156        image = image.astype(np.float32)157        image = image / 255158        label = np.ceil(label / 255)159        for i in range(6):160            image[:, :, i] -= self.mean[i]161        for i in range(6):162            image[:, :, i] /= self.std[i]163 164        return [image, label]165 166 167class GaussianNoise(object):168    def __init__(self, std=0.05):169        '''170        :param mean: global mean computed from dataset171        :param std: global std computed from dataset172        '''173        self.std = std174 175    def __call__(self, image, label):176        noise = np.random.normal(loc=0, scale=self.std, size=image.shape)177        image = image + noise.astype(np.float32)178        return [image, label]179 180 181class ToTensor(object):182    '''183    This class converts the data to tensor so that it can be processed by PyTorch184    '''185 186    def __init__(self, scale=1):187        '''188        :param scale: set this parameter according to the output scale189        '''190        self.scale = scale191 192    def __call__(self, image, label):193        if self.scale != 1:194            h, w = label.shape[:2]195            image = cv2.resize(image, (int(w), int(h)))196            label = cv2.resize(label, (int(w / self.scale), int(h / self.scale)), \197                               interpolation=cv2.INTER_NEAREST)198        image = image[:, :, ::-1].copy()  # .copy() is to solve "torch does not support negative index"199        image = image.transpose((2, 0, 1))200        image_tensor = torch.from_numpy(image)201        label_tensor = torch.LongTensor(np.array(label, dtype=np.int)).unsqueeze(dim=0)202 203        return [image_tensor, label_tensor]204 205 206class Compose(object):207    """208    Composes several transforms together.209    """210 211    def __init__(self, transforms):212        self.transforms = transforms213 214    def __call__(self, *args):215        for t in self.transforms:216            args = t(*args)217        return args218