init

lib/datasets/utils/dataset_colors.py

@@ -0,0 +1,77 @@
+import random
+from functools import lru_cache
+import numpy as np
+
+dataset_color_dict = {
+    "potsdam" : [[1], [2], [3], [4], [5]],
+    "vaihingen" : [[255, 255, 0], [0, 255, 0], [0, 255, 255], [0, 0, 255], [255, 255, 255]],
+    "deepglobe" : [[255,255,255], [0,0,255], [0,255,0],[255,0,255], [255,255,0], [0,255,255]],
+    "fbp" : [[i+1] for i in range(24)],
+    "loveda" : [[i+2, i+2, i+2] for i in range(6)],
+    "isaid" : [[i+1] for i in range(15)],
+    "pastis-mm" : [[i+1] for i in range(18)],
+    "dynamic-mm" : [[i] for i in range(7)], 
+    "c2seg-ab" : [[i+1] for i in range(13)],
+    "flood3i": [[i+1] for i in range(9)],
+    "jl16-mm": [[i] for i in range(16)],
+    "flair-mm": [[i+1] for i in range(18)],
+    "dfc20": [[i+1] for i in range(10)]
+}
+
+
+modal_norm_dict = {
+    'hr' : {
+        'div' : 255.,
+        'mean' : [0.485, 0.456, 0.406],
+        'std' : [0.229, 0.224, 0.225]
+    },
+    'anno' : {
+        'div' : 255.,
+        'mean' : [0.485, 0.456, 0.406],
+        'std' : [0.229, 0.224, 0.225]
+    },
+    's2' : {
+        'div' : 1.,
+        'mean' : [884.29673756, 1144.16202635, 1297.47289228, 1624.90992062, 2194.6423161, 2422.21248945, 2517.76053101, 2581.64687018, 2368.51236873, 1805.06846033],
+        'std' : [1155.15170768, 1183.6292542, 1368.11351514, 1370.265037, 1355.55390699, 1416.51487101, 1474.78900051, 1439.3086061, 1455.52084939, 1343.48379601]
+    },
+    's1' : {
+        'div' : 1.,
+        'mean' : [-12.54847273, -20.19237134],
+        'std' : [5.25697717, 5.91150917]
+    },
+}
+
+@lru_cache()
+def get_painter_color_map_list(num_locations = 300):
+    
+    num_sep_per_channel = int(num_locations ** (1 / 3)) + 1  # 19
+    separation_per_channel = 256 // num_sep_per_channel
+
+    color_list = []
+    for location in range(num_locations):
+        num_seq_r = location // num_sep_per_channel ** 2
+        num_seq_g = (location % num_sep_per_channel ** 2) // num_sep_per_channel
+        num_seq_b = location % num_sep_per_channel
+        assert (num_seq_r <= num_sep_per_channel) and (num_seq_g <= num_sep_per_channel) \
+               and (num_seq_b <= num_sep_per_channel)
+
+        R = 255 - num_seq_r * separation_per_channel
+        G = 255 - num_seq_g * separation_per_channel
+        B = 255 - num_seq_b * separation_per_channel
+        assert (R < 256) and (G < 256) and (B < 256)
+        assert (R >= 0) and (G >= 0) and (B >= 0)
+        assert (R, G, B) not in color_list
+
+        color_list.append((R, G, B))
+
+    return color_list
+
+
+def get_real_random_color_list(num_locations):
+    random_color_list = np.random.randint(0, 256, (num_locations, 3))
+    while np.sum(random_color_list) == 0:
+        print('random_color_list is 0!')
+        random_color_list = np.random.randint(0, 256, (num_locations, 3))
+    random_color_list = random_color_list.tolist()
+    return random_color_list # [:num_locations]

lib/datasets/utils/formatting.py

@@ -0,0 +1,65 @@
+from collections.abc import Sequence
+import mmcv
+import numpy as np
+import torch
+
+
+def to_tensor(data):
+    """Convert objects of various python types to :obj:`torch.Tensor`.
+
+    Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
+    :class:`Sequence`, :class:`int` and :class:`float`.
+
+    Args:
+        data (torch.Tensor | numpy.ndarray | Sequence | int | float): Data to
+            be converted.
+    """
+
+    if isinstance(data, torch.Tensor):
+        return data
+    elif isinstance(data, np.ndarray):
+        return torch.from_numpy(data)
+    elif isinstance(data, Sequence) and not mmcv.is_str(data):
+        return torch.tensor(data)
+    elif isinstance(data, int):
+        return torch.LongTensor([data])
+    elif isinstance(data, float):
+        return torch.FloatTensor([data])
+    else:
+        raise TypeError(f'type {type(data)} cannot be converted to tensor.')
+
+
+class ToTensor(object):
+    """Convert some sample to :obj:`torch.Tensor` by given keys.
+
+    Args:
+        keys (Sequence[str]): Keys that need to be converted to Tensor.
+    """
+
+    def __init__(self, keys):
+        self.keys = keys
+
+    def __call__(self, sample):
+        """Call function to convert data in sample to :obj:`torch.Tensor`.
+
+        Args:
+            sample (Sample): sample data contains the data to convert.
+
+        Returns:
+            dict: The result dict contains the data converted
+                to :obj:`torch.Tensor`.
+        """
+
+        for key in self.keys:
+            if isinstance(sample[key], list):
+                for i in range(len(sample[key])):
+                    sample[key][i] = to_tensor(sample[key][i])
+            else:
+                sample[key] = to_tensor(sample[key])
+        return sample
+
+    def __repr__(self):
+        return self.__class__.__name__ + f'(keys={self.keys})'
+
+
+

lib/datasets/utils/masking_generator.py

@@ -0,0 +1,84 @@
+import random
+import math
+import numpy as np
+
+class MaskingGenerator:
+    def __init__(
+            self, input_size, patch_size, mask_ratio=0.5, min_num_patches=4, max_num_patches=None,
+            min_aspect=0.3, max_aspect=None):
+        if not isinstance(input_size, list):
+            input_size = [input_size,] * 2
+        self.height = input_size[0] // patch_size
+        self.width = input_size[1] // patch_size
+
+        self.num_patches = self.height * self.width
+        self.num_masking_patches = int(self.num_patches * mask_ratio)
+
+        self.min_num_patches = min_num_patches
+        self.max_num_patches = self.num_masking_patches if max_num_patches is None else max_num_patches
+
+        max_aspect = max_aspect or 1 / min_aspect
+        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))
+
+    def __repr__(self):
+        repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
+            self.height, self.width, self.min_num_patches, self.max_num_patches,
+            self.num_masking_patches, self.log_aspect_ratio[0], self.log_aspect_ratio[1])
+        return repr_str
+
+    def get_shape(self):
+        return self.height, self.width
+
+    def _mask(self, mask, max_mask_patches):
+        delta = 0
+        for attempt in range(10):
+            target_area = random.uniform(self.min_num_patches, max_mask_patches)
+            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+            h = int(round(math.sqrt(target_area * aspect_ratio)))
+            w = int(round(math.sqrt(target_area / aspect_ratio)))
+            if w < self.width and h < self.height:
+                top = random.randint(0, self.height - h)
+                left = random.randint(0, self.width - w)
+
+                num_masked = mask[top: top + h, left: left + w].sum()
+                # Overlap
+                if 0 < h * w - num_masked <= max_mask_patches:
+                    for i in range(top, top + h):
+                        for j in range(left, left + w):
+                            if mask[i, j] == 0:
+                                mask[i, j] = 1
+                                delta += 1
+
+                if delta > 0:
+                    break
+        return delta
+
+    def __call__(self):
+        mask = np.zeros(shape=self.get_shape(), dtype=np.int32)
+        mask_count = 0
+        while mask_count < self.num_masking_patches:
+            max_mask_patches = self.num_masking_patches - mask_count
+            max_mask_patches = min(max_mask_patches, self.max_num_patches)
+
+            delta = self._mask(mask, max_mask_patches)
+            if delta == 0:
+                break
+            else:
+                mask_count += delta
+
+        # maintain a fix number {self.num_masking_patches}
+        if mask_count > self.num_masking_patches:
+            delta = mask_count - self.num_masking_patches
+            mask_x, mask_y = mask.nonzero()
+            to_vis = np.random.choice(mask_x.shape[0], delta, replace=False)
+            mask[mask_x[to_vis], mask_y[to_vis]] = 0
+
+        elif mask_count < self.num_masking_patches:
+            delta = self.num_masking_patches - mask_count
+            mask_x, mask_y = (mask == 0).nonzero()
+            to_mask = np.random.choice(mask_x.shape[0], delta, replace=False)
+            mask[mask_x[to_mask], mask_y[to_mask]] = 1
+
+        assert mask.sum() == self.num_masking_patches, f"mask: {mask}, mask count {mask.sum()}"
+
+        return mask

lib/datasets/utils/pair_trainsforms.py

@@ -0,0 +1,532 @@
+import math
+import numbers
+import random
+import warnings
+from collections.abc import Sequence
+from typing import List, Optional, Tuple
+import numpy as np
+
+import torch
+from torch import Tensor
+import torchvision.transforms as transforms
+from skimage import io
+
+
+try:
+    import accimage
+except ImportError:
+    accimage = None
+
+import torchvision.transforms.functional as F
+from torchvision.transforms.functional import _interpolation_modes_from_int, InterpolationMode
+from PIL import Image, ImageFilter, ImageOps
+
+from .dataset_colors import modal_norm_dict
+
+__all__ = [
+    "Compose",
+    "ToTensor",
+    "Normalize",
+    "RandomHorizontalFlip",
+    "RandomResizedCrop",
+]
+
+
+
+class Compose(transforms.Compose):
+    """Composes several transforms together. This transform does not support torchscript.
+    Please, see the note below.
+    Args:
+        transforms (list of ``Transform`` objects): list of transforms to compose.
+    """
+
+    def __init__(self, transforms):
+        super().__init__(transforms)
+
+    def __call__(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None):
+        # i = 0
+        for t in self.transforms:
+            # i = i+1
+            # print(f'dataset_name:{dataset_name}')
+            # print(f'step:{i}')
+            # print(f'hr_img shape:{hr_img.shape}')
+            # print(f's2_img shape:{s2_img.shape}')
+            # print(f's1_img shape:{s1_img.shape}')
+            # print(f'tgt shape:{tgt.shape}')
+            
+            hr_img, s2_img, s1_img, tgt = t(dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=interpolation1, interpolation2=interpolation2)
+        return hr_img, s2_img, s1_img, tgt
+
+
+class ToTensor(transforms.ToTensor):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor. This transform does not support torchscript.
+    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
+    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
+    if the PIL Image belongs to one of the modes (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1)
+    or if the numpy.ndarray has dtype = np.uint8
+    In the other cases, tensors are returned without scaling.
+    .. note::
+        Because the input image is scaled to [0.0, 1.0], this transformation should not be used when
+        transforming target image masks. See the `references`_ for implementing the transforms for image masks.
+    .. _references: https://github.com/pytorch/vision/tree/main/references/segmentation
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __call__(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+        Returns:
+            Tensor: Converted image.
+        """
+        
+
+        # print(f'hr dtype:{hr_img.dtype}')
+        # print(f's2_img dtype:{s2_img.dtype}')
+        # print(f's1_img dtype:{s1_img.dtype}')
+        # print(f'tgt dtype:{tgt.dtype}')
+        if dataset_name == 'dynamic-mm' or dataset_name == 'guizhou-mm':
+            hr_img = hr_img.astype(np.int32)[:3,:,:]
+            hr_img = hr_img[::-1,:,:].copy()
+        else:
+            hr_img = hr_img.astype(np.int32)
+        tgt = tgt.astype(np.uint8)
+        s1_img = s1_img.astype(np.float32)
+        s2_img = s2_img.astype(np.int16)
+        
+        return torch.tensor(hr_img), torch.tensor(s2_img), torch.tensor(s1_img),torch.tensor(tgt)
+
+
+class Normalize(transforms.Normalize):
+    """Normalize a tensor image with mean and standard deviation.
+    This transform does not support PIL Image.
+    Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n``
+    channels, this transform will normalize each channel of the input
+    ``torch.*Tensor`` i.e.,
+    ``output[channel] = (input[channel] - mean[channel]) / std[channel]``
+    .. note::
+        This transform acts out of place, i.e., it does not mutate the input tensor.
+    Args:
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channel.
+        inplace(bool,optional): Bool to make this operation in-place.
+    """
+
+    def __init__(self, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], inplace=False):
+        super().__init__(mean, std, inplace)
+
+    def forward(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            tensor (Tensor): Tensor image to be normalized.
+        Returns:
+            Tensor: Normalized Tensor image.
+        """
+        # TODO 查询对应的mean和std
+        
+        # 处理一些mean和std
+        if dataset_name == 'dynamic-mm':
+            hr_std = [1008.4052, 760.9586, 631.4754]
+            hr_mean = [1085.2941, 944.2718, 689.2493]
+            hr_div = 1.
+        else:
+            hr_mean = modal_norm_dict['hr']['mean']
+            hr_std = modal_norm_dict['hr']['std']
+            hr_div = modal_norm_dict['hr']['div']            
+
+        if dataset_name == 'l8activefire':
+        # if False:
+            s2_mean = modal_norm_dict['l8']['mean']
+            s2_std = modal_norm_dict['l8']['std']
+            s2_div = modal_norm_dict['l8']['div']
+        else:
+            s2_mean = modal_norm_dict['s2']['mean']
+            s2_std = modal_norm_dict['s2']['std']
+            s2_div = modal_norm_dict['s2']['div']
+            
+        s1_mean = modal_norm_dict['s1']['mean']
+        s1_std = modal_norm_dict['s1']['std']
+        s1_div = modal_norm_dict['s1']['div']
+
+        anno_mean = [0.485, 0.456, 0.406]
+        anno_std = [0.229, 0.224, 0.225]
+        ann_div = 255.
+
+        # 存在问题：时间序列这样处理是否会出错
+
+        #mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device)
+        #std = torch.as_tensor(std, dtype=dtype, device=tensor.device)
+        # print(s2_img.shape)
+        # import pdb; pdb.set_trace()
+        # print(s2_img)
+        try:
+            ch, ts, h, w = s2_img.shape
+        except:
+            print(f's2: {s2_img.shape}, s1: {s1_img.shape}')
+        s2_img = s2_img.view(ch, ts*h, w)
+        s2_img = self.normalize(s2_img.type(torch.float32), s2_mean, s2_std, self.inplace)
+        s2_img = s2_img.view(ch, ts, h, w)
+
+        ch, ts, h, w = s1_img.shape
+        s1_img = s1_img.view(ch, ts*h, w)
+        s1_img = self.normalize(s1_img.type(torch.float32), s1_mean, s1_std, self.inplace)
+        s1_img = s1_img.view(ch, ts, h, w)
+
+        # import pdb; pdb.set_trace()
+        # print(s2_img.shape, s2_img[:,0,:,:])
+        # print(s1_img.shape, s1_img[:,0,:,:])
+        # print(hr_mean, hr_std, hr_div)
+        return self.normalize(hr_img.type(torch.float32).div_(hr_div), hr_mean, hr_std, self.inplace), \
+            s2_img, \
+            s1_img, \
+            self.normalize(tgt.type(torch.float32).div_(ann_div) , anno_mean, anno_std, self.inplace)
+
+    def normalize(self, tensor, mean, std, inplace):
+        dtype = tensor.dtype
+        mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device)
+        std = torch.as_tensor(std, dtype=dtype, device=tensor.device)
+        if (std == 0).any():
+            raise ValueError(f"std evaluated to zero after conversion to {dtype}, leading to division by zero.")
+        mean = mean.view(-1, 1, 1)
+        std = std.view(-1, 1, 1)
+        # print(f'tensor shape: {tensor.shape}')
+        # print(f'mean shape: {mean.shape}')
+        # print(f'std shape: {std.shape}')
+        return tensor.sub_(mean).div_(std) 
+
+
+class RandomResizedCrop(transforms.RandomResizedCrop):
+    """Crop a random portion of image and resize it to a given size.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading dimensions
+    A crop of the original image is made: the crop has a random area (H * W)
+    and a random aspect ratio. This crop is finally resized to the given
+    size. This is popularly used to train the Inception networks.
+    Args:
+        size (int or sequence): expected output size of the crop, for each edge. If size is an
+            int instead of sequence like (h, w), a square output size ``(size, size)`` is
+            made. If provided a sequence of length 1, it will be interpreted as (size[0], size[0]).
+            .. note::
+                In torchscript mode size as single int is not supported, use a sequence of length 1: ``[size, ]``.
+        scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
+            before resizing. The scale is defined with respect to the area of the original image.
+        ratio (tuple of float): lower and upper bounds for the random aspect ratio of the crop, before
+            resizing.
+        interpolation (InterpolationMode): Desired interpolation enum defined by
+            :class:`torchvision.transforms.InterpolationMode`. Default is ``InterpolationMode.BILINEAR``.
+            If input is Tensor, only ``InterpolationMode.NEAREST``, ``InterpolationMode.BILINEAR`` and
+            ``InterpolationMode.BICUBIC`` are supported.
+            For backward compatibility integer values (e.g. ``PIL.Image[.Resampling].NEAREST``) are still accepted,
+            but deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum.
+    """
+
+    def __init__(
+        self,
+        size,
+        scale=(0.08, 1.0),
+        ratio=(3.0 / 4.0, 4.0 / 3.0),
+        interpolation=InterpolationMode.BILINEAR,
+        mode='small'
+    ):
+        super().__init__(size, scale=scale, ratio=ratio, interpolation=interpolation)
+        self.cnt=0
+        self.mode = mode
+
+    def forward(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None, mode='small'):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be cropped and resized.
+        Returns:
+            PIL Image or Tensor: Randomly cropped and resized image.
+        """
+        i, j, h, w = self.get_params(s2_img, self.scale, self.ratio)
+        size_hr = hr_img.shape[-1]
+        size_s2 = s2_img.shape[-1]
+        size_anno = tgt.shape[-1]
+        # 映射到其他模态
+        ratio_s2_hr = size_s2 / size_hr
+        i_hr = int(i / ratio_s2_hr)
+        j_hr = int(j / ratio_s2_hr)
+        h_hr = int(h / ratio_s2_hr)
+        w_hr = int(w / ratio_s2_hr)
+
+        ratio_s2_anno = size_s2 / size_anno
+        i_anno = int(i / ratio_s2_anno)
+        j_anno = int(j / ratio_s2_anno)
+        h_anno = int(h / ratio_s2_anno)
+        w_anno = int(w / ratio_s2_anno)
+
+        if interpolation1 == 'nearest':
+            interpolation1 = InterpolationMode.NEAREST
+        else:
+            interpolation1 = InterpolationMode.BICUBIC
+        if interpolation2 == 'nearest':
+            interpolation2 = InterpolationMode.NEAREST
+        else:
+            interpolation2 = InterpolationMode.BICUBIC
+        # import pdb;pdb.set_trace()
+        if self.scale[0]>0.99 and self.scale[0]<1.0:
+            if self.mode=='small':
+                resized_s2_img = F.resize(s2_img, (16,16), interpolation=InterpolationMode.BICUBIC)
+                resized_hr_img = F.resize(hr_img, (512, 512), interpolation=InterpolationMode.BICUBIC)
+                resized_s1_img = F.resize(s1_img, (16,16), interpolation=InterpolationMode.BICUBIC)
+                resized_tgt = F.resize(tgt, (512,512), interpolation=InterpolationMode.NEAREST)
+            else:
+                resized_s2_img = F.resize(s2_img, (64,64), interpolation=InterpolationMode.BICUBIC)
+                resized_hr_img = F.resize(hr_img, (2048, 2048), interpolation=InterpolationMode.BICUBIC)
+                resized_s1_img = F.resize(s1_img, (64,64), interpolation=InterpolationMode.BICUBIC)
+                resized_tgt = F.resize(tgt, (2048,2048), interpolation=InterpolationMode.NEAREST)
+            return resized_hr_img, resized_s2_img, resized_s1_img, resized_tgt
+
+        if self.mode=='small':
+            resized_s2_img =  F.resized_crop(s2_img, i, j, h, w, (16, 16), InterpolationMode.BICUBIC)
+            resized_hr_img = F.resized_crop(hr_img, i_hr, j_hr, h_hr, w_hr, (512, 512), InterpolationMode.BICUBIC)
+            resized_s1_img = F.resized_crop(s1_img, i, j, h, w, (16, 16), InterpolationMode.BICUBIC)
+            resized_tgt = F.resized_crop(tgt, i_anno, j_anno, h_anno, w_anno, (512, 512), InterpolationMode.NEAREST)
+        else:
+            resized_s2_img =  F.resized_crop(s2_img, i, j, h, w, (512, 512), InterpolationMode.BICUBIC)
+            resized_hr_img = F.resized_crop(hr_img, i_hr, j_hr, h_hr, w_hr, (2048,2048), InterpolationMode.BICUBIC)
+            resized_s1_img = F.resized_crop(s1_img, i, j, h, w, (512, 512), InterpolationMode.BICUBIC)
+            resized_tgt = F.resized_crop(tgt, i_anno, j_anno, h_anno, w_anno, (2048, 2048), InterpolationMode.NEAREST)
+
+        # import pdb; pdb.set_trace()
+        # 将resize后的结果保存为concat的img
+        # self.cnt = self.cnt+1
+        # from torchvision.utils import save_image
+        # save_hr = resized_hr_img[:3, :, :] / resized_hr_img[:3, :, :].max()
+        # save_s2 = resized_s2_img[:3,0,:,:] / resized_s2_img[:3,0,:,:].max()
+        # print(f'{save_hr.shape}, {save_s2.shape}')
+        # save_image(save_s2, f'FoundationModel/debug/output2/resized_s2_{self.cnt}.png')
+        # save_image(save_hr, f'FoundationModel/debug/output2/resized_hr_{self.cnt}.png')
+        
+        return resized_hr_img, resized_s2_img, resized_s1_img, resized_tgt
+        
+class RandomResizedCropComb(transforms.RandomResizedCrop):
+    def __init__(
+        self,
+        size,
+        scale=(0.08, 1.0),
+        ratio=(3.0 / 4.0, 4.0 / 3.0),
+        interpolation=InterpolationMode.BILINEAR,
+    ):
+        super().__init__(size, scale=scale, ratio=ratio, interpolation=interpolation)
+        self.cnt=0
+
+    def forward(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be cropped and resized.
+        Returns:
+            PIL Image or Tensor: Randomly cropped and resized image.
+        """
+        i, j, h, w = self.get_params(s2_img, self.scale, self.ratio)
+        # print(f'i, j, h, w: {i, j, h, w}')
+        # print(f's2_img shape: {s2_img.shape}')
+        size_hr = hr_img.shape[-1]
+        size_s2 = s2_img.shape[-1]
+        size_anno = tgt.shape[-1]
+        # 映射到其他模态
+        ratio_s2_hr = size_s2 / size_hr
+        i_hr = int(i / ratio_s2_hr)
+        j_hr = int(j / ratio_s2_hr)
+        h_hr = int(h / ratio_s2_hr)
+        w_hr = int(w / ratio_s2_hr)
+
+        ratio_s2_anno = size_s2 / size_anno
+        i_anno = int(i / ratio_s2_anno)
+        j_anno = int(j / ratio_s2_anno)
+        h_anno = int(h / ratio_s2_anno)
+        w_anno = int(w / ratio_s2_anno)
+
+        if interpolation1 == 'nearest':
+            interpolation1 = InterpolationMode.NEAREST
+        else:
+            interpolation1 = InterpolationMode.BICUBIC
+        if interpolation2 == 'nearest':
+            interpolation2 = InterpolationMode.NEAREST
+        else:
+            interpolation2 = InterpolationMode.BICUBIC
+        
+        resized_s2_img =  F.resized_crop(s2_img, i, j, h, w, (32, 16), InterpolationMode.BICUBIC)
+        resized_hr_img = F.resized_crop(hr_img, i_hr, j_hr, h_hr, w_hr, (1024, 512), InterpolationMode.BICUBIC)
+        resized_s1_img = F.resized_crop(s1_img, i, j, h, w, (32, 16), InterpolationMode.BICUBIC)
+        resized_tgt = F.resized_crop(tgt, i_anno, j_anno, h_anno, w_anno, (1024, 512), InterpolationMode.NEAREST)
+  
+        return resized_hr_img, resized_s2_img, resized_s1_img, resized_tgt
+        
+
+class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
+    """Horizontally flip the given image randomly with a given probability.
+    If the image is torch Tensor, it is expected
+    to have [..., H, W] shape, where ... means an arbitrary number of leading
+    dimensions
+    Args:
+        p (float): probability of the image being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        super().__init__(p=p)
+
+    def forward(self, dataset_name, hr_img, s2_img, s1_img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            img (PIL Image or Tensor): Image to be flipped.
+        Returns:
+            PIL Image or Tensor: Randomly flipped image.
+        """
+        if torch.rand(1) < self.p:
+            return F.hflip(hr_img), F.hflip(s2_img), F.hflip(s1_img), F.hflip(tgt)
+        return hr_img, s2_img, s1_img, tgt
+
+
+class RandomApply(transforms.RandomApply):
+    """Apply randomly a list of transformations with a given probability.
+    .. note::
+        In order to script the transformation, please use ``torch.nn.ModuleList`` as input instead of list/tuple of
+        transforms as shown below:
+        >>> transforms = transforms.RandomApply(torch.nn.ModuleList([
+        >>>     transforms.ColorJitter(),
+        >>> ]), p=0.3)
+        >>> scripted_transforms = torch.jit.script(transforms)
+        Make sure to use only scriptable transformations, i.e. that work with ``torch.Tensor``, does not require
+        `lambda` functions or ``PIL.Image``.
+    Args:
+        transforms (sequence or torch.nn.Module): list of transformations
+        p (float): probability
+    """
+
+    def __init__(self, transforms, p=0.5):
+        super().__init__(transforms, p=p)
+
+    def forward(self, img, tgt, interpolation1=None, interpolation2=None):
+        if self.p < torch.rand(1):
+            return img, tgt
+        for t in self.transforms:
+            img, tgt = t(img, tgt)
+        return img, tgt
+
+class ColorJitter(transforms.ColorJitter):
+    """Randomly change the brightness, contrast, saturation and hue of an image.
+    If the image is torch Tensor, it is expected
+    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, mode "1", "I", "F" and modes with transparency (alpha channel) are not supported.
+    Args:
+        brightness (float or tuple of float (min, max)): How much to jitter brightness.
+            brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
+            or the given [min, max]. Should be non negative numbers.
+        contrast (float or tuple of float (min, max)): How much to jitter contrast.
+            contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]
+            or the given [min, max]. Should be non negative numbers.
+        saturation (float or tuple of float (min, max)): How much to jitter saturation.
+            saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]
+            or the given [min, max]. Should be non negative numbers.
+        hue (float or tuple of float (min, max)): How much to jitter hue.
+            hue_factor is chosen uniformly from [-hue, hue] or the given [min, max].
+            Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5.
+            To jitter hue, the pixel values of the input image has to be non-negative for conversion to HSV space;
+            thus it does not work if you normalize your image to an interval with negative values,
+            or use an interpolation that generates negative values before using this function.
+    """
+
+    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
+        super().__init__(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)
+
+    def forward(self, img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            img (PIL Image or Tensor): Input image.
+        Returns:
+            PIL Image or Tensor: Color jittered image.
+        """
+        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = self.get_params(
+            self.brightness, self.contrast, self.saturation, self.hue
+        )
+
+        for fn_id in fn_idx:
+            if fn_id == 0 and brightness_factor is not None:
+                img = F.adjust_brightness(img, brightness_factor)
+            elif fn_id == 1 and contrast_factor is not None:
+                img = F.adjust_contrast(img, contrast_factor)
+            elif fn_id == 2 and saturation_factor is not None:
+                img = F.adjust_saturation(img, saturation_factor)
+            elif fn_id == 3 and hue_factor is not None:
+                img = F.adjust_hue(img, hue_factor)
+        return img, tgt
+
+
+class RandomErasing(transforms.RandomErasing):
+    """Randomly selects a rectangle region in a torch.Tensor image and erases its pixels.
+    This transform does not support PIL Image.
+    'Random Erasing Data Augmentation' by Zhong et al. See https://arxiv.org/abs/1708.04896
+    Args:
+         p: probability that the random erasing operation will be performed.
+         scale: range of proportion of erased area against input image.
+         ratio: range of aspect ratio of erased area.
+         value: erasing value. Default is 0. If a single int, it is used to
+            erase all pixels. If a tuple of length 3, it is used to erase
+            R, G, B channels respectively.
+            If a str of 'random', erasing each pixel with random values.
+         inplace: boolean to make this transform inplace. Default set to False.
+    Returns:
+        Erased Image.
+    Example:
+        >>> transform = transforms.Compose([
+        >>>   transforms.RandomHorizontalFlip(),
+        >>>   transforms.PILToTensor(),
+        >>>   transforms.ConvertImageDtype(torch.float),
+        >>>   transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        >>>   transforms.RandomErasing(),
+        >>> ])
+    """
+
+    def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False):
+        super().__init__(p=p, scale=scale, ratio=ratio, value=value, inplace=inplace)
+
+    def forward(self, img, tgt, interpolation1=None, interpolation2=None):
+        """
+        Args:
+            img (Tensor): Tensor image to be erased.
+        Returns:
+            img (Tensor): Erased Tensor image.
+        """
+        if torch.rand(1) < self.p:
+
+            # cast self.value to script acceptable type
+            if isinstance(self.value, (int, float)):
+                value = [self.value]
+            elif isinstance(self.value, str):
+                value = None
+            elif isinstance(self.value, tuple):
+                value = list(self.value)
+            else:
+                value = self.value
+
+            if value is not None and not (len(value) in (1, img.shape[-3])):
+                raise ValueError(
+                    "If value is a sequence, it should have either a single value or "
+                    f"{img.shape[-3]} (number of input channels)"
+                )
+
+            x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=value)
+            return F.erase(img, x, y, h, w, v, self.inplace), tgt
+        return img, tgt
+
+
+
+class GaussianBlur(object):
+    """Gaussian blur augmentation from SimCLR: https://arxiv.org/abs/2002.05709"""
+
+    def __init__(self, sigma=[.1, 2.]):
+        self.sigma = sigma
+
+    def __call__(self, img, tgt, interpolation1=None, interpolation2=None):
+        sigma = random.uniform(self.sigma[0], self.sigma[1])
+        img = img.filter(ImageFilter.GaussianBlur(radius=sigma))
+        return img, tgt
+
+    def __repr__(self) -> str:
+        s = f"{self.__class__.__name__}( sigma={self.sigma})"
+        return s
+