123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109 |
- import os
- import copy
- import random
- from numbers import Number
- from functools import partial
- from operator import methodcaller
- from collections.abc import Sequence
- import numpy as np
- import cv2
- import imghdr
- from PIL import Image
- from joblib import load
- import paddlers
- import paddlers.transforms.indices as indices
- import paddlers.transforms.satellites as satellites
- from .functions import (
- normalize, horizontal_flip, permute, vertical_flip, center_crop, is_poly,
- horizontal_flip_poly, horizontal_flip_rle, vertical_flip_poly,
- vertical_flip_rle, crop_poly, crop_rle, expand_poly, expand_rle,
- resize_poly, resize_rle, dehaze, select_bands, to_intensity, to_uint8,
- img_flip, img_simple_rotate, decode_seg_mask, calc_hr_shape,
- match_by_regression, match_histograms)
- __all__ = [
- "Compose",
- "DecodeImg",
- "Resize",
- "RandomResize",
- "ResizeByShort",
- "RandomResizeByShort",
- "ResizeByLong",
- "RandomHorizontalFlip",
- "RandomVerticalFlip",
- "Normalize",
- "CenterCrop",
- "RandomCrop",
- "RandomScaleAspect",
- "RandomExpand",
- "Pad",
- "MixupImage",
- "RandomDistort",
- "RandomBlur",
- "RandomSwap",
- "Dehaze",
- "ReduceDim",
- "SelectBand",
- "RandomFlipOrRotate",
- "ReloadMask",
- "AppendIndex",
- "MatchRadiance",
- "ArrangeRestorer",
- "ArrangeSegmenter",
- "ArrangeChangeDetector",
- "ArrangeClassifier",
- "ArrangeDetector",
- ]
- interp_dict = {
- 'NEAREST': cv2.INTER_NEAREST,
- 'LINEAR': cv2.INTER_LINEAR,
- 'CUBIC': cv2.INTER_CUBIC,
- 'AREA': cv2.INTER_AREA,
- 'LANCZOS4': cv2.INTER_LANCZOS4
- }
- class Compose(object):
- """
- Apply a series of data augmentation strategies to the input.
- All input images should be in Height-Width-Channel ([H, W, C]) format.
- Args:
- transforms (list[paddlers.transforms.Transform]): List of data preprocess or
- augmentation operators.
- Raises:
- TypeError: Invalid type of transforms.
- ValueError: Invalid length of transforms.
- """
- def __init__(self, transforms):
- super(Compose, self).__init__()
- if not isinstance(transforms, list):
- raise TypeError(
- "Type of transforms is invalid. Must be a list, but received is {}."
- .format(type(transforms)))
- if len(transforms) < 1:
- raise ValueError(
- "Length of transforms must not be less than 1, but received is {}."
- .format(len(transforms)))
- transforms = copy.deepcopy(transforms)
- self.arrange = self._pick_arrange(transforms)
- self.transforms = transforms
- def __call__(self, sample):
- """
- This is equivalent to sequentially calling compose_obj.apply_transforms()
- and compose_obj.arrange_outputs().
- """
- sample = self.apply_transforms(sample)
- sample = self.arrange_outputs(sample)
- return sample
- def apply_transforms(self, sample):
- for op in self.transforms:
-
- if isinstance(op, (paddlers.transforms.BatchRandomResize,
- paddlers.transforms.BatchRandomResizeByShort,
- MixupImage)):
- continue
- sample = op(sample)
- return sample
- def arrange_outputs(self, sample):
- if self.arrange is not None:
- sample = self.arrange(sample)
- return sample
- def _pick_arrange(self, transforms):
- arrange = None
- for idx, op in enumerate(transforms):
- if isinstance(op, Arrange):
- if idx != len(transforms) - 1:
- raise ValueError(
- "Arrange operator must be placed at the end of the list."
- )
- arrange = transforms.pop(idx)
- return arrange
- class Transform(object):
- """
- Parent class of all data augmentation operators.
- """
- def __init__(self):
- pass
- def apply_im(self, image):
- return image
- def apply_mask(self, mask):
- return mask
- def apply_bbox(self, bbox):
- return bbox
- def apply_segm(self, segms):
- return segms
- def apply(self, sample):
- if 'image' in sample:
- sample['image'] = self.apply_im(sample['image'])
- else:
- sample['image'] = self.apply_im(sample['image_t1'])
- sample['image2'] = self.apply_im(sample['image_t2'])
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'])
- if 'gt_bbox' in sample:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'])
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(self.apply_mask, sample['aux_masks']))
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- def __call__(self, sample):
- if isinstance(sample, Sequence):
- sample = [self.apply(s) for s in sample]
- else:
- sample = self.apply(sample)
- return sample
- class DecodeImg(Transform):
- """
- Decode image(s) in input.
-
- Args:
- to_rgb (bool, optional): If True, convert input image(s) from BGR format to
- RGB format. Defaults to True.
- to_uint8 (bool, optional): If True, quantize and convert decoded image(s) to
- uint8 type. Defaults to True.
- decode_bgr (bool, optional): If True, automatically interpret a non-geo image
- (e.g., jpeg images) as a BGR image. Defaults to True.
- decode_sar (bool, optional): If True, automatically interpret a single-channel
- geo image (e.g. geotiff images) as a SAR image, set this argument to
- True. Defaults to True.
- read_geo_info (bool, optional): If True, read geographical information from
- the image. Deafults to False.
- use_stretch (bool, optional): Whether to apply 2% linear stretch. Valid only if
- `to_uint8` is True. Defaults to False.
- """
- def __init__(self,
- to_rgb=True,
- to_uint8=True,
- decode_bgr=True,
- decode_sar=True,
- read_geo_info=False,
- use_stretch=False):
- super(DecodeImg, self).__init__()
- self.to_rgb = to_rgb
- self.to_uint8 = to_uint8
- self.decode_bgr = decode_bgr
- self.decode_sar = decode_sar
- self.read_geo_info = read_geo_info
- self.use_stretch = use_stretch
- def read_img(self, img_path):
- img_format = imghdr.what(img_path)
- name, ext = os.path.splitext(img_path)
- geo_trans, geo_proj = None, None
- if img_format == 'tiff' or ext == '.img':
- try:
- import gdal
- except:
- try:
- from osgeo import gdal
- except ImportError:
- raise ImportError(
- "Failed to import gdal! Please install GDAL library according to the document."
- )
- dataset = gdal.Open(img_path)
- if dataset == None:
- raise IOError('Cannot open', img_path)
- im_data = dataset.ReadAsArray()
- if im_data.ndim == 2 and self.decode_sar:
- im_data = to_intensity(im_data)
- im_data = im_data[:, :, np.newaxis]
- else:
- if im_data.ndim == 3:
- im_data = im_data.transpose((1, 2, 0))
- if self.read_geo_info:
- geo_trans = dataset.GetGeoTransform()
- geo_proj = dataset.GetProjection()
- elif img_format in ['jpeg', 'bmp', 'png', 'jpg']:
- if self.decode_bgr:
- im_data = cv2.imread(img_path, cv2.IMREAD_ANYDEPTH |
- cv2.IMREAD_ANYCOLOR | cv2.IMREAD_COLOR)
- else:
- im_data = cv2.imread(img_path, cv2.IMREAD_ANYDEPTH |
- cv2.IMREAD_ANYCOLOR)
- elif ext == '.npy':
- im_data = np.load(img_path)
- else:
- raise TypeError("Image format {} is not supported!".format(ext))
- if self.read_geo_info:
- return im_data, geo_trans, geo_proj
- else:
- return im_data
- def apply_im(self, im_path):
- if isinstance(im_path, str):
- try:
- data = self.read_img(im_path)
- except:
- raise ValueError("Cannot read the image file {}!".format(
- im_path))
- if self.read_geo_info:
- image, geo_trans, geo_proj = data
- geo_info_dict = {'geo_trans': geo_trans, 'geo_proj': geo_proj}
- else:
- image = data
- else:
- image = im_path
- if self.to_rgb and image.shape[-1] == 3:
- image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
- if self.to_uint8:
- image = to_uint8(image, stretch=self.use_stretch)
- if self.read_geo_info:
- return image, geo_info_dict
- else:
- return image
- def apply_mask(self, mask):
- try:
- mask = np.asarray(Image.open(mask))
- except:
- raise ValueError("Cannot read the mask file {}!".format(mask))
- if len(mask.shape) != 2:
- raise ValueError(
- "Mask should be a 1-channel image, but recevied is a {}-channel image.".
- format(mask.shape[2]))
- return mask
- def apply(self, sample):
- """
- Args:
- sample (dict): Input sample.
- Returns:
- dict: Sample with decoded images.
- """
- if 'image' in sample:
- if self.read_geo_info:
- image, geo_info_dict = self.apply_im(sample['image'])
- sample['image'] = image
- sample['geo_info_dict'] = geo_info_dict
- else:
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- if self.read_geo_info:
- image2, geo_info_dict2 = self.apply_im(sample['image2'])
- sample['image2'] = image2
- sample['geo_info_dict2'] = geo_info_dict2
- else:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'image_t1' in sample and not 'image' in sample:
- if not ('image_t2' in sample and 'image2' not in sample):
- raise ValueError
- if self.read_geo_info:
- image, geo_info_dict = self.apply_im(sample['image_t1'])
- sample['image'] = image
- sample['geo_info_dict'] = geo_info_dict
- else:
- sample['image'] = self.apply_im(sample['image_t1'])
- if self.read_geo_info:
- image2, geo_info_dict2 = self.apply_im(sample['image_t2'])
- sample['image2'] = image2
- sample['geo_info_dict2'] = geo_info_dict2
- else:
- sample['image2'] = self.apply_im(sample['image_t2'])
- if 'mask' in sample:
- sample['mask_ori'] = copy.deepcopy(sample['mask'])
- sample['mask'] = self.apply_mask(sample['mask'])
- im_height, im_width, _ = sample['image'].shape
- se_height, se_width = sample['mask'].shape
- if im_height != se_height or im_width != se_width:
- raise ValueError(
- "The height or width of the image is not same as the mask.")
- if 'aux_masks' in sample:
- sample['aux_masks_ori'] = copy.deepcopy(sample['aux_masks'])
- sample['aux_masks'] = list(
- map(self.apply_mask, sample['aux_masks']))
-
- if 'target' in sample:
- if self.read_geo_info:
- target, geo_info_dict = self.apply_im(sample['target'])
- sample['target'] = target
- sample['geo_info_dict_tar'] = geo_info_dict
- else:
- sample['target'] = self.apply_im(sample['target'])
- sample['im_shape'] = np.array(
- sample['image'].shape[:2], dtype=np.float32)
- sample['scale_factor'] = np.array([1., 1.], dtype=np.float32)
- return sample
- class Resize(Transform):
- """
- Resize input.
- - If `target_size` is an int, resize the image(s) to (`target_size`, `target_size`).
- - If `target_size` is a list or tuple, resize the image(s) to `target_size`.
- Attention: If `interp` is 'RANDOM', the interpolation method will be chosen randomly.
- Args:
- target_size (int | list[int] | tuple[int]): Target size. If it is an integer, the
- target height and width will be both set to `target_size`. Otherwise,
- `target_size` represents [target height, target width].
- interp (str, optional): Interpolation method for resizing image(s). One of
- {'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}.
- Defaults to 'LINEAR'.
- keep_ratio (bool, optional): If True, the scaling factor of width and height will
- be set to same value, and height/width of the resized image will be not
- greater than the target width/height. Defaults to False.
- Raises:
- TypeError: Invalid type of target_size.
- ValueError: Invalid interpolation method.
- """
- def __init__(self, target_size, interp='LINEAR', keep_ratio=False):
- super(Resize, self).__init__()
- if not (interp == "RANDOM" or interp in interp_dict):
- raise ValueError("`interp` should be one of {}.".format(
- interp_dict.keys()))
- if isinstance(target_size, int):
- target_size = (target_size, target_size)
- else:
- if not (isinstance(target_size,
- (list, tuple)) and len(target_size) == 2):
- raise TypeError(
- "`target_size` should be an int or a list of length 2, but received {}.".
- format(target_size))
-
- self.target_size = target_size
- self.interp = interp
- self.keep_ratio = keep_ratio
- def apply_im(self, image, interp, target_size):
- flag = image.shape[2] == 1
- image = cv2.resize(image, target_size, interpolation=interp)
- if flag:
- image = image[:, :, np.newaxis]
- return image
- def apply_mask(self, mask, target_size):
- mask = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
- return mask
- def apply_bbox(self, bbox, scale, target_size):
- im_scale_x, im_scale_y = scale
- bbox[:, 0::2] *= im_scale_x
- bbox[:, 1::2] *= im_scale_y
- bbox[:, 0::2] = np.clip(bbox[:, 0::2], 0, target_size[0])
- bbox[:, 1::2] = np.clip(bbox[:, 1::2], 0, target_size[1])
- return bbox
- def apply_segm(self, segms, im_size, scale):
- im_h, im_w = im_size
- im_scale_x, im_scale_y = scale
- resized_segms = []
- for segm in segms:
- if is_poly(segm):
-
- resized_segms.append([
- resize_poly(poly, im_scale_x, im_scale_y) for poly in segm
- ])
- else:
-
- resized_segms.append(
- resize_rle(segm, im_h, im_w, im_scale_x, im_scale_y))
- return resized_segms
- def apply(self, sample):
- if self.interp == "RANDOM":
- interp = random.choice(list(interp_dict.values()))
- else:
- interp = interp_dict[self.interp]
- im_h, im_w = sample['image'].shape[:2]
- im_scale_y = self.target_size[0] / im_h
- im_scale_x = self.target_size[1] / im_w
- target_size = (self.target_size[1], self.target_size[0])
- if self.keep_ratio:
- scale = min(im_scale_y, im_scale_x)
- target_w = int(round(im_w * scale))
- target_h = int(round(im_h * scale))
- target_size = (target_w, target_h)
- im_scale_y = target_h / im_h
- im_scale_x = target_w / im_w
- sample['image'] = self.apply_im(sample['image'], interp, target_size)
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], interp,
- target_size)
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'], target_size)
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(partial(
- self.apply_mask, target_size=target_size),
- sample['aux_masks']))
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(
- sample['gt_bbox'], [im_scale_x, im_scale_y], target_size)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(
- sample['gt_poly'], [im_h, im_w], [im_scale_x, im_scale_y])
- if 'target' in sample:
- if 'sr_factor' in sample:
-
- sample['target'] = self.apply_im(
- sample['target'], interp,
- calc_hr_shape(target_size, sample['sr_factor']))
- else:
-
- sample['target'] = self.apply_im(sample['target'], interp,
- target_size)
- sample['im_shape'] = np.asarray(
- sample['image'].shape[:2], dtype=np.float32)
- if 'scale_factor' in sample:
- scale_factor = sample['scale_factor']
- sample['scale_factor'] = np.asarray(
- [scale_factor[0] * im_scale_y, scale_factor[1] * im_scale_x],
- dtype=np.float32)
- return sample
- class RandomResize(Transform):
- """
- Resize input to random sizes.
- Attention: If `interp` is 'RANDOM', the interpolation method will be chosen randomly.
- Args:
- target_sizes (list[int] | list[list|tuple] | tuple[list|tuple]):
- Multiple target sizes, each of which should be int, list, or tuple.
- interp (str, optional): Interpolation method for resizing image(s). One of
- {'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}.
- Defaults to 'LINEAR'.
- Raises:
- TypeError: Invalid type of `target_size`.
- ValueError: Invalid interpolation method.
- """
- def __init__(self, target_sizes, interp='LINEAR'):
- super(RandomResize, self).__init__()
- if not (interp == "RANDOM" or interp in interp_dict):
- raise ValueError("`interp` should be one of {}.".format(
- interp_dict.keys()))
- self.interp = interp
- assert isinstance(target_sizes, list), \
- "`target_size` must be a list."
- for i, item in enumerate(target_sizes):
- if isinstance(item, int):
- target_sizes[i] = (item, item)
- self.target_size = target_sizes
- def apply(self, sample):
- height, width = random.choice(self.target_size)
- resizer = Resize((height, width), interp=self.interp)
- sample = resizer(sample)
- return sample
- class ResizeByShort(Transform):
- """
- Resize input while keeping the aspect ratio.
- Attention: If `interp` is 'RANDOM', the interpolation method will be chosen randomly.
- Args:
- short_size (int): Target size of the shorter side of the image(s).
- max_size (int, optional): Upper bound of longer side of the image(s). If
- `max_size` is -1, no upper bound will be applied. Defaults to -1.
- interp (str, optional): Interpolation method for resizing image(s). One of
- {'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}.
- Defaults to 'LINEAR'.
- Raises:
- ValueError: Invalid interpolation method.
- """
- def __init__(self, short_size=256, max_size=-1, interp='LINEAR'):
- if not (interp == "RANDOM" or interp in interp_dict):
- raise ValueError("`interp` should be one of {}".format(
- interp_dict.keys()))
- super(ResizeByShort, self).__init__()
- self.short_size = short_size
- self.max_size = max_size
- self.interp = interp
- def apply(self, sample):
- im_h, im_w = sample['image'].shape[:2]
- im_short_size = min(im_h, im_w)
- im_long_size = max(im_h, im_w)
- scale = float(self.short_size) / float(im_short_size)
- if 0 < self.max_size < np.round(scale * im_long_size):
- scale = float(self.max_size) / float(im_long_size)
- target_w = int(round(im_w * scale))
- target_h = int(round(im_h * scale))
- sample = Resize(
- target_size=(target_h, target_w), interp=self.interp)(sample)
- return sample
- class RandomResizeByShort(Transform):
- """
- Resize input to random sizes while keeping the aspect ratio.
- Attention: If `interp` is 'RANDOM', the interpolation method will be chosen randomly.
- Args:
- short_sizes (list[int]): Target size of the shorter side of the image(s).
- max_size (int, optional): Upper bound of longer side of the image(s).
- If `max_size` is -1, no upper bound will be applied. Defaults to -1.
- interp (str, optional): Interpolation method for resizing image(s). One of
- {'NEAREST', 'LINEAR', 'CUBIC', 'AREA', 'LANCZOS4', 'RANDOM'}.
- Defaults to 'LINEAR'.
- Raises:
- TypeError: Invalid type of `target_size`.
- ValueError: Invalid interpolation method.
- See Also:
- ResizeByShort: Resize image(s) in input while keeping the aspect ratio.
- """
- def __init__(self, short_sizes, max_size=-1, interp='LINEAR'):
- super(RandomResizeByShort, self).__init__()
- if not (interp == "RANDOM" or interp in interp_dict):
- raise ValueError("`interp` should be one of {}".format(
- interp_dict.keys()))
- self.interp = interp
- assert isinstance(short_sizes, list), \
- "`short_sizes` must be a list."
- self.short_sizes = short_sizes
- self.max_size = max_size
- def apply(self, sample):
- short_size = random.choice(self.short_sizes)
- resizer = ResizeByShort(
- short_size=short_size, max_size=self.max_size, interp=self.interp)
- sample = resizer(sample)
- return sample
- class ResizeByLong(Transform):
- def __init__(self, long_size=256, interp='LINEAR'):
- super(ResizeByLong, self).__init__()
- self.long_size = long_size
- self.interp = interp
- def apply(self, sample):
- im_h, im_w = sample['image'].shape[:2]
- im_long_size = max(im_h, im_w)
- scale = float(self.long_size) / float(im_long_size)
- target_h = int(round(im_h * scale))
- target_w = int(round(im_w * scale))
- sample = Resize(
- target_size=(target_h, target_w), interp=self.interp)(sample)
- return sample
- class RandomFlipOrRotate(Transform):
- """
- Flip or Rotate an image in different directions with a certain probability.
- Args:
- probs (list[float]): Probabilities of performing flipping and rotation.
- Default: [0.35,0.25].
- probsf (list[float]): Probabilities of 5 flipping modes (horizontal,
- vertical, both horizontal and vertical, diagonal, anti-diagonal).
- Default: [0.3, 0.3, 0.2, 0.1, 0.1].
- probsr (list[float]): Probabilities of 3 rotation modes (90°, 180°, 270°
- clockwise). Default: [0.25, 0.5, 0.25].
- Examples:
- from paddlers import transforms as T
- # Define operators for data augmentation
- train_transforms = T.Compose([
- T.DecodeImg(),
- T.RandomFlipOrRotate(
- probs = [0.3, 0.2] # p=0.3 to flip the image,p=0.2 to rotate the image,p=0.5 to keep the image unchanged.
- probsf = [0.3, 0.25, 0, 0, 0] # p=0.3 and p=0.25 to perform horizontal and vertical flipping; probility of no-flipping is 0.45.
- probsr = [0, 0.65, 0]), # p=0.65 to rotate the image by 180°; probility of no-rotation is 0.35.
- T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
- ])
- """
- def __init__(self,
- probs=[0.35, 0.25],
- probsf=[0.3, 0.3, 0.2, 0.1, 0.1],
- probsr=[0.25, 0.5, 0.25]):
- super(RandomFlipOrRotate, self).__init__()
-
- self.probs = [probs[0], probs[0] + probs[1]]
- self.probsf = self.get_probs_range(probsf)
- self.probsr = self.get_probs_range(probsr)
- def apply_im(self, image, mode_id, flip_mode=True):
- if flip_mode:
- image = img_flip(image, mode_id)
- else:
- image = img_simple_rotate(image, mode_id)
- return image
- def apply_mask(self, mask, mode_id, flip_mode=True):
- if flip_mode:
- mask = img_flip(mask, mode_id)
- else:
- mask = img_simple_rotate(mask, mode_id)
- return mask
- def apply_bbox(self, bbox, mode_id, flip_mode=True):
- raise TypeError(
- "Currently, RandomFlipOrRotate is not available for object detection tasks."
- )
- def apply_segm(self, bbox, mode_id, flip_mode=True):
- raise TypeError(
- "Currently, RandomFlipOrRotate is not available for object detection tasks."
- )
- def get_probs_range(self, probs):
- """
- Change list of probabilities into cumulative probability intervals.
- Args:
- probs (list[float]): Probabilities of different modes, shape: [n].
- Returns:
- list[list]: Probability intervals, shape: [n, 2].
- """
- ps = []
- last_prob = 0
- for prob in probs:
- p_s = last_prob
- cur_prob = prob / sum(probs)
- last_prob += cur_prob
- p_e = last_prob
- ps.append([p_s, p_e])
- return ps
- def judge_probs_range(self, p, probs):
- """
- Judge whether the value of `p` falls within the given probability interval.
- Args:
- p (float): Value between 0 and 1.
- probs (list[list]): Probability intervals, shape: [n, 2].
- Returns:
- int: Interval where the input probability falls into.
- """
- for id, id_range in enumerate(probs):
- if p > id_range[0] and p < id_range[1]:
- return id
- return -1
- def apply(self, sample):
- p_m = random.random()
- if p_m < self.probs[0]:
- mode_p = random.random()
- mode_id = self.judge_probs_range(mode_p, self.probsf)
- sample['image'] = self.apply_im(sample['image'], mode_id, True)
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], mode_id,
- True)
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'], mode_id, True)
- if 'aux_masks' in sample:
- sample['aux_masks'] = [
- self.apply_mask(aux_mask, mode_id, True)
- for aux_mask in sample['aux_masks']
- ]
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], mode_id,
- True)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(sample['gt_poly'], mode_id,
- True)
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'], mode_id,
- True)
- elif p_m < self.probs[1]:
- mode_p = random.random()
- mode_id = self.judge_probs_range(mode_p, self.probsr)
- sample['image'] = self.apply_im(sample['image'], mode_id, False)
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], mode_id,
- False)
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'], mode_id, False)
- if 'aux_masks' in sample:
- sample['aux_masks'] = [
- self.apply_mask(aux_mask, mode_id, False)
- for aux_mask in sample['aux_masks']
- ]
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], mode_id,
- False)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(sample['gt_poly'], mode_id,
- False)
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'], mode_id,
- False)
- return sample
- class RandomHorizontalFlip(Transform):
- """
- Randomly flip the input horizontally.
- Args:
- prob (float, optional): Probability of flipping the input. Defaults to .5.
- """
- def __init__(self, prob=0.5):
- super(RandomHorizontalFlip, self).__init__()
- self.prob = prob
- def apply_im(self, image):
- image = horizontal_flip(image)
- return image
- def apply_mask(self, mask):
- mask = horizontal_flip(mask)
- return mask
- def apply_bbox(self, bbox, width):
- oldx1 = bbox[:, 0].copy()
- oldx2 = bbox[:, 2].copy()
- bbox[:, 0] = width - oldx2
- bbox[:, 2] = width - oldx1
- return bbox
- def apply_segm(self, segms, height, width):
- flipped_segms = []
- for segm in segms:
- if is_poly(segm):
-
- flipped_segms.append(
- [horizontal_flip_poly(poly, width) for poly in segm])
- else:
-
- flipped_segms.append(horizontal_flip_rle(segm, height, width))
- return flipped_segms
- def apply(self, sample):
- if random.random() < self.prob:
- im_h, im_w = sample['image'].shape[:2]
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'])
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(self.apply_mask, sample['aux_masks']))
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], im_w)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(sample['gt_poly'], im_h,
- im_w)
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class RandomVerticalFlip(Transform):
- """
- Randomly flip the input vertically.
- Args:
- prob (float, optional): Probability of flipping the input. Defaults to .5.
- """
- def __init__(self, prob=0.5):
- super(RandomVerticalFlip, self).__init__()
- self.prob = prob
- def apply_im(self, image):
- image = vertical_flip(image)
- return image
- def apply_mask(self, mask):
- mask = vertical_flip(mask)
- return mask
- def apply_bbox(self, bbox, height):
- oldy1 = bbox[:, 1].copy()
- oldy2 = bbox[:, 3].copy()
- bbox[:, 0] = height - oldy2
- bbox[:, 2] = height - oldy1
- return bbox
- def apply_segm(self, segms, height, width):
- flipped_segms = []
- for segm in segms:
- if is_poly(segm):
-
- flipped_segms.append(
- [vertical_flip_poly(poly, height) for poly in segm])
- else:
-
- flipped_segms.append(vertical_flip_rle(segm, height, width))
- return flipped_segms
- def apply(self, sample):
- if random.random() < self.prob:
- im_h, im_w = sample['image'].shape[:2]
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'])
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(self.apply_mask, sample['aux_masks']))
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], im_h)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(sample['gt_poly'], im_h,
- im_w)
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class Normalize(Transform):
- """
- Apply normalization to the input image(s). The normalization steps are:
- 1. im = (im - min_value) * 1 / (max_value - min_value)
- 2. im = im - mean
- 3. im = im / std
- Args:
- mean (list[float] | tuple[float], optional): Mean of input image(s).
- Defaults to [0.485, 0.456, 0.406].
- std (list[float] | tuple[float], optional): Standard deviation of input
- image(s). Defaults to [0.229, 0.224, 0.225].
- min_val (list[float] | tuple[float], optional): Minimum value of input
- image(s). If None, use 0 for all channels. Defaults to None.
- max_val (list[float] | tuple[float], optional): Maximum value of input
- image(s). If None, use 255. for all channels. Defaults to None.
- apply_to_tar (bool, optional): Whether to apply transformation to the target
- image. Defaults to True.
- """
- def __init__(self,
- mean=[0.485, 0.456, 0.406],
- std=[0.229, 0.224, 0.225],
- min_val=None,
- max_val=None,
- apply_to_tar=True):
- super(Normalize, self).__init__()
- channel = len(mean)
- if min_val is None:
- min_val = [0] * channel
- if max_val is None:
- max_val = [255.] * channel
- from functools import reduce
- if reduce(lambda x, y: x * y, std) == 0:
- raise ValueError(
- "`std` should not contain 0, but received is {}.".format(std))
- if reduce(lambda x, y: x * y,
- [a - b for a, b in zip(max_val, min_val)]) == 0:
- raise ValueError(
- "(`max_val` - `min_val`) should not contain 0, but received is {}.".
- format((np.asarray(max_val) - np.asarray(min_val)).tolist()))
- self.mean = mean
- self.std = std
- self.min_val = min_val
- self.max_val = max_val
- self.apply_to_tar = apply_to_tar
- def apply_im(self, image):
- image = image.astype(np.float32)
- mean = np.asarray(
- self.mean, dtype=np.float32)[np.newaxis, np.newaxis, :]
- std = np.asarray(self.std, dtype=np.float32)[np.newaxis, np.newaxis, :]
- image = normalize(image, mean, std, self.min_val, self.max_val)
- return image
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'target' in sample and self.apply_to_tar:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class CenterCrop(Transform):
- """
- Crop the input image(s) at the center.
- 1. Locate the center of the image.
- 2. Crop the image.
- Args:
- crop_size (int, optional): Target size of the cropped image(s).
- Defaults to 224.
- """
- def __init__(self, crop_size=224):
- super(CenterCrop, self).__init__()
- self.crop_size = crop_size
- def apply_im(self, image):
- image = center_crop(image, self.crop_size)
- return image
- def apply_mask(self, mask):
- mask = center_crop(mask, self.crop_size)
- return mask
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'])
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(self.apply_mask, sample['aux_masks']))
- if 'target' in sample:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class RandomCrop(Transform):
- """
- Randomly crop the input.
- 1. Compute the height and width of cropped area according to `aspect_ratio` and
- `scaling`.
- 2. Locate the upper left corner of cropped area randomly.
- 3. Crop the image(s).
- 4. Resize the cropped area to `crop_size` x `crop_size`.
- Args:
- crop_size (int | list[int] | tuple[int]): Target size of the cropped area. If
- None, the cropped area will not be resized. Defaults to None.
- aspect_ratio (list[float], optional): Aspect ratio of cropped region in
- [min, max] format. Defaults to [.5, 2.].
- thresholds (list[float], optional): Iou thresholds to decide a valid bbox
- crop. Defaults to [.0, .1, .3, .5, .7, .9].
- scaling (list[float], optional): Ratio between the cropped region and the
- original image in [min, max] format. Defaults to [.3, 1.].
- num_attempts (int, optional): Max number of tries before giving up.
- Defaults to 50.
- allow_no_crop (bool, optional): Whether returning without doing crop is
- allowed. Defaults to True.
- cover_all_box (bool, optional): Whether to ensure all bboxes be covered in
- the final crop. Defaults to False.
- """
- def __init__(self,
- crop_size=None,
- aspect_ratio=[.5, 2.],
- thresholds=[.0, .1, .3, .5, .7, .9],
- scaling=[.3, 1.],
- num_attempts=50,
- allow_no_crop=True,
- cover_all_box=False):
- super(RandomCrop, self).__init__()
- self.crop_size = crop_size
- self.aspect_ratio = aspect_ratio
- self.thresholds = thresholds
- self.scaling = scaling
- self.num_attempts = num_attempts
- self.allow_no_crop = allow_no_crop
- self.cover_all_box = cover_all_box
- def _generate_crop_info(self, sample):
- im_h, im_w = sample['image'].shape[:2]
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- thresholds = self.thresholds
- if self.allow_no_crop:
- thresholds.append('no_crop')
- np.random.shuffle(thresholds)
- for thresh in thresholds:
- if thresh == 'no_crop':
- return None
- for i in range(self.num_attempts):
- crop_box = self._get_crop_box(im_h, im_w)
- if crop_box is None:
- continue
- iou = self._iou_matrix(
- sample['gt_bbox'],
- np.array(
- [crop_box], dtype=np.float32))
- if iou.max() < thresh:
- continue
- if self.cover_all_box and iou.min() < thresh:
- continue
- cropped_box, valid_ids = self._crop_box_with_center_constraint(
- sample['gt_bbox'], np.array(
- crop_box, dtype=np.float32))
- if valid_ids.size > 0:
- return crop_box, cropped_box, valid_ids
- else:
- for i in range(self.num_attempts):
- crop_box = self._get_crop_box(im_h, im_w)
- if crop_box is None:
- continue
- return crop_box, None, None
- return None
- def _get_crop_box(self, im_h, im_w):
- scale = np.random.uniform(*self.scaling)
- if self.aspect_ratio is not None:
- min_ar, max_ar = self.aspect_ratio
- aspect_ratio = np.random.uniform(
- max(min_ar, scale**2), min(max_ar, scale**-2))
- h_scale = scale / np.sqrt(aspect_ratio)
- w_scale = scale * np.sqrt(aspect_ratio)
- else:
- h_scale = np.random.uniform(*self.scaling)
- w_scale = np.random.uniform(*self.scaling)
- crop_h = im_h * h_scale
- crop_w = im_w * w_scale
- if self.aspect_ratio is None:
- if crop_h / crop_w < 0.5 or crop_h / crop_w > 2.0:
- return None
- crop_h = int(crop_h)
- crop_w = int(crop_w)
- crop_y = np.random.randint(0, im_h - crop_h)
- crop_x = np.random.randint(0, im_w - crop_w)
- return [crop_x, crop_y, crop_x + crop_w, crop_y + crop_h]
- def _iou_matrix(self, a, b):
- tl_i = np.maximum(a[:, np.newaxis, :2], b[:, :2])
- br_i = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
- area_i = np.prod(br_i - tl_i, axis=2) * (tl_i < br_i).all(axis=2)
- area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
- area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
- area_o = (area_a[:, np.newaxis] + area_b - area_i)
- return area_i / (area_o + 1e-10)
- def _crop_box_with_center_constraint(self, box, crop):
- cropped_box = box.copy()
- cropped_box[:, :2] = np.maximum(box[:, :2], crop[:2])
- cropped_box[:, 2:] = np.minimum(box[:, 2:], crop[2:])
- cropped_box[:, :2] -= crop[:2]
- cropped_box[:, 2:] -= crop[:2]
- centers = (box[:, :2] + box[:, 2:]) / 2
- valid = np.logical_and(crop[:2] <= centers,
- centers < crop[2:]).all(axis=1)
- valid = np.logical_and(
- valid, (cropped_box[:, :2] < cropped_box[:, 2:]).all(axis=1))
- return cropped_box, np.where(valid)[0]
- def _crop_segm(self, segms, valid_ids, crop, height, width):
- crop_segms = []
- for id in valid_ids:
- segm = segms[id]
- if is_poly(segm):
-
- crop_segms.append(crop_poly(segm, crop))
- else:
-
- crop_segms.append(crop_rle(segm, crop, height, width))
- return crop_segms
- def apply_im(self, image, crop):
- x1, y1, x2, y2 = crop
- return image[y1:y2, x1:x2, :]
- def apply_mask(self, mask, crop):
- x1, y1, x2, y2 = crop
- return mask[y1:y2, x1:x2, ...]
- def apply(self, sample):
- crop_info = self._generate_crop_info(sample)
- if crop_info is not None:
- crop_box, cropped_box, valid_ids = crop_info
- im_h, im_w = sample['image'].shape[:2]
- sample['image'] = self.apply_im(sample['image'], crop_box)
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], crop_box)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- crop_polys = self._crop_segm(
- sample['gt_poly'],
- valid_ids,
- np.array(
- crop_box, dtype=np.int64),
- im_h,
- im_w)
- if [] in crop_polys:
- delete_id = list()
- valid_polys = list()
- for idx, poly in enumerate(crop_polys):
- if not crop_poly:
- delete_id.append(idx)
- else:
- valid_polys.append(poly)
- valid_ids = np.delete(valid_ids, delete_id)
- if not valid_polys:
- return sample
- sample['gt_poly'] = valid_polys
- else:
- sample['gt_poly'] = crop_polys
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = np.take(cropped_box, valid_ids, axis=0)
- sample['gt_class'] = np.take(
- sample['gt_class'], valid_ids, axis=0)
- if 'gt_score' in sample:
- sample['gt_score'] = np.take(
- sample['gt_score'], valid_ids, axis=0)
- if 'is_crowd' in sample:
- sample['is_crowd'] = np.take(
- sample['is_crowd'], valid_ids, axis=0)
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'], crop_box)
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(partial(
- self.apply_mask, crop=crop_box),
- sample['aux_masks']))
- if 'target' in sample:
- if 'sr_factor' in sample:
- sample['target'] = self.apply_im(
- sample['target'],
- calc_hr_shape(crop_box, sample['sr_factor']))
- else:
- sample['target'] = self.apply_im(sample['image'], crop_box)
- if self.crop_size is not None:
- sample = Resize(self.crop_size)(sample)
- return sample
- class RandomScaleAspect(Transform):
- """
- Crop input image(s) and resize back to original sizes.
- Args:
- min_scale (float): Minimum ratio between the cropped region and the original
- image. If 0, image(s) will not be cropped. Defaults to .5.
- aspect_ratio (float): Aspect ratio of cropped region. Defaults to .33.
- """
- def __init__(self, min_scale=0.5, aspect_ratio=0.33):
- super(RandomScaleAspect, self).__init__()
- self.min_scale = min_scale
- self.aspect_ratio = aspect_ratio
- def apply(self, sample):
- if self.min_scale != 0 and self.aspect_ratio != 0:
- img_height, img_width = sample['image'].shape[:2]
- sample = RandomCrop(
- crop_size=(img_height, img_width),
- aspect_ratio=[self.aspect_ratio, 1. / self.aspect_ratio],
- scaling=[self.min_scale, 1.],
- num_attempts=10,
- allow_no_crop=False)(sample)
- return sample
- class RandomExpand(Transform):
- """
- Randomly expand the input by padding according to random offsets.
- Args:
- upper_ratio (float, optional): Maximum ratio to which the original image
- is expanded. Defaults to 4..
- prob (float, optional): Probability of apply expanding. Defaults to .5.
- im_padding_value (list[float] | tuple[float], optional): RGB filling value
- for the image. Defaults to (127.5, 127.5, 127.5).
- label_padding_value (int, optional): Filling value for the mask.
- Defaults to 255.
- See Also:
- paddlers.transforms.Pad
- """
- def __init__(self,
- upper_ratio=4.,
- prob=.5,
- im_padding_value=127.5,
- label_padding_value=255):
- super(RandomExpand, self).__init__()
- assert upper_ratio > 1.01, "`upper_ratio` must be larger than 1.01."
- self.upper_ratio = upper_ratio
- self.prob = prob
- assert isinstance(im_padding_value, (Number, Sequence)), \
- "Value to fill must be either float or sequence."
- self.im_padding_value = im_padding_value
- self.label_padding_value = label_padding_value
- def apply(self, sample):
- if random.random() < self.prob:
- im_h, im_w = sample['image'].shape[:2]
- ratio = np.random.uniform(1., self.upper_ratio)
- h = int(im_h * ratio)
- w = int(im_w * ratio)
- if h > im_h and w > im_w:
- y = np.random.randint(0, h - im_h)
- x = np.random.randint(0, w - im_w)
- target_size = (h, w)
- offsets = (x, y)
- sample = Pad(
- target_size=target_size,
- pad_mode=-1,
- offsets=offsets,
- im_padding_value=self.im_padding_value,
- label_padding_value=self.label_padding_value)(sample)
- return sample
- class Pad(Transform):
- def __init__(self,
- target_size=None,
- pad_mode=0,
- offsets=None,
- im_padding_value=127.5,
- label_padding_value=255,
- size_divisor=32):
- """
- Pad image to a specified size or multiple of `size_divisor`.
- Args:
- target_size (list[int] | tuple[int], optional): Image target size, if None, pad to
- multiple of size_divisor. Defaults to None.
- pad_mode (int, optional): Pad mode. Currently only four modes are supported:
- [-1, 0, 1, 2]. if -1, use specified offsets. If 0, only pad to right and bottom
- If 1, pad according to center. If 2, only pad left and top. Defaults to 0.
- offsets (list[int]|None, optional): Padding offsets. Defaults to None.
- im_padding_value (list[float] | tuple[float]): RGB value of padded area.
- Defaults to (127.5, 127.5, 127.5).
- label_padding_value (int, optional): Filling value for the mask.
- Defaults to 255.
- size_divisor (int): Image width and height after padding will be a multiple of
- `size_divisor`.
- """
- super(Pad, self).__init__()
- if isinstance(target_size, (list, tuple)):
- if len(target_size) != 2:
- raise ValueError(
- "`target_size` should contain 2 elements, but it is {}.".
- format(target_size))
- if isinstance(target_size, int):
- target_size = [target_size] * 2
- assert pad_mode in [
- -1, 0, 1, 2
- ], "Currently only four modes are supported: [-1, 0, 1, 2]."
- if pad_mode == -1:
- assert offsets, "if `pad_mode` is -1, `offsets` should not be None."
- self.target_size = target_size
- self.size_divisor = size_divisor
- self.pad_mode = pad_mode
- self.offsets = offsets
- self.im_padding_value = im_padding_value
- self.label_padding_value = label_padding_value
- def apply_im(self, image, offsets, target_size):
- x, y = offsets
- h, w = target_size
- im_h, im_w, channel = image.shape[:3]
- canvas = np.ones((h, w, channel), dtype=np.float32)
- canvas *= np.array(self.im_padding_value, dtype=np.float32)
- canvas[y:y + im_h, x:x + im_w, :] = image.astype(np.float32)
- return canvas
- def apply_mask(self, mask, offsets, target_size):
- x, y = offsets
- im_h, im_w = mask.shape[:2]
- h, w = target_size
- canvas = np.ones((h, w), dtype=np.float32)
- canvas *= np.array(self.label_padding_value, dtype=np.float32)
- canvas[y:y + im_h, x:x + im_w] = mask.astype(np.float32)
- return canvas
- def apply_bbox(self, bbox, offsets):
- return bbox + np.array(offsets * 2, dtype=np.float32)
- def apply_segm(self, segms, offsets, im_size, size):
- x, y = offsets
- height, width = im_size
- h, w = size
- expanded_segms = []
- for segm in segms:
- if is_poly(segm):
-
- expanded_segms.append(
- [expand_poly(poly, x, y) for poly in segm])
- else:
-
- expanded_segms.append(
- expand_rle(segm, x, y, height, width, h, w))
- return expanded_segms
- def _get_offsets(self, im_h, im_w, h, w):
- if self.pad_mode == -1:
- offsets = self.offsets
- elif self.pad_mode == 0:
- offsets = [0, 0]
- elif self.pad_mode == 1:
- offsets = [(w - im_w) // 2, (h - im_h) // 2]
- else:
- offsets = [w - im_w, h - im_h]
- return offsets
- def apply(self, sample):
- im_h, im_w = sample['image'].shape[:2]
- if self.target_size:
- h, w = self.target_size
- assert (
- im_h <= h and im_w <= w
- ), 'target size ({}, {}) cannot be less than image size ({}, {})'\
- .format(h, w, im_h, im_w)
- else:
- h = (np.ceil(im_h / self.size_divisor) *
- self.size_divisor).astype(int)
- w = (np.ceil(im_w / self.size_divisor) *
- self.size_divisor).astype(int)
- if h == im_h and w == im_w:
- return sample
- offsets = self._get_offsets(im_h, im_w, h, w)
- sample['image'] = self.apply_im(sample['image'], offsets, (h, w))
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], offsets, (h, w))
- if 'mask' in sample:
- sample['mask'] = self.apply_mask(sample['mask'], offsets, (h, w))
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(partial(
- self.apply_mask, offsets=offsets, target_size=(h, w)),
- sample['aux_masks']))
- if 'gt_bbox' in sample and len(sample['gt_bbox']) > 0:
- sample['gt_bbox'] = self.apply_bbox(sample['gt_bbox'], offsets)
- if 'gt_poly' in sample and len(sample['gt_poly']) > 0:
- sample['gt_poly'] = self.apply_segm(
- sample['gt_poly'], offsets, im_size=[im_h, im_w], size=[h, w])
- if 'target' in sample:
- if 'sr_factor' in sample:
- hr_shape = calc_hr_shape((h, w), sample['sr_factor'])
- hr_offsets = self._get_offsets(*sample['target'].shape[:2],
- *hr_shape)
- sample['target'] = self.apply_im(sample['target'], hr_offsets,
- hr_shape)
- else:
- sample['target'] = self.apply_im(sample['target'], offsets,
- (h, w))
- return sample
- class MixupImage(Transform):
- def __init__(self, alpha=1.5, beta=1.5, mixup_epoch=-1):
- """
- Mixup two images and their gt_bbbox/gt_score.
- Args:
- alpha (float, optional): Alpha parameter of beta distribution.
- Defaults to 1.5.
- beta (float, optional): Beta parameter of beta distribution.
- Defaults to 1.5.
- """
- super(MixupImage, self).__init__()
- if alpha <= 0.0:
- raise ValueError("`alpha` should be positive in MixupImage.")
- if beta <= 0.0:
- raise ValueError("`beta` should be positive in MixupImage.")
- self.alpha = alpha
- self.beta = beta
- self.mixup_epoch = mixup_epoch
- def apply_im(self, image1, image2, factor):
- h = max(image1.shape[0], image2.shape[0])
- w = max(image1.shape[1], image2.shape[1])
- img = np.zeros((h, w, image1.shape[2]), 'float32')
- img[:image1.shape[0], :image1.shape[1], :] = \
- image1.astype('float32') * factor
- img[:image2.shape[0], :image2.shape[1], :] += \
- image2.astype('float32') * (1.0 - factor)
- return img.astype('uint8')
- def __call__(self, sample):
- if not isinstance(sample, Sequence):
- return sample
- assert len(sample) == 2, 'mixup need two samples'
- factor = np.random.beta(self.alpha, self.beta)
- factor = max(0.0, min(1.0, factor))
- if factor >= 1.0:
- return sample[0]
- if factor <= 0.0:
- return sample[1]
- image = self.apply_im(sample[0]['image'], sample[1]['image'], factor)
- result = copy.deepcopy(sample[0])
- result['image'] = image
-
- if 'gt_bbox' in sample[0]:
- gt_bbox1 = sample[0]['gt_bbox']
- gt_bbox2 = sample[1]['gt_bbox']
- gt_bbox = np.concatenate((gt_bbox1, gt_bbox2), axis=0)
- result['gt_bbox'] = gt_bbox
- if 'gt_poly' in sample[0]:
- gt_poly1 = sample[0]['gt_poly']
- gt_poly2 = sample[1]['gt_poly']
- gt_poly = gt_poly1 + gt_poly2
- result['gt_poly'] = gt_poly
- if 'gt_class' in sample[0]:
- gt_class1 = sample[0]['gt_class']
- gt_class2 = sample[1]['gt_class']
- gt_class = np.concatenate((gt_class1, gt_class2), axis=0)
- result['gt_class'] = gt_class
- gt_score1 = np.ones_like(sample[0]['gt_class'])
- gt_score2 = np.ones_like(sample[1]['gt_class'])
- gt_score = np.concatenate(
- (gt_score1 * factor, gt_score2 * (1. - factor)), axis=0)
- result['gt_score'] = gt_score
- if 'is_crowd' in sample[0]:
- is_crowd1 = sample[0]['is_crowd']
- is_crowd2 = sample[1]['is_crowd']
- is_crowd = np.concatenate((is_crowd1, is_crowd2), axis=0)
- result['is_crowd'] = is_crowd
- if 'difficult' in sample[0]:
- is_difficult1 = sample[0]['difficult']
- is_difficult2 = sample[1]['difficult']
- is_difficult = np.concatenate(
- (is_difficult1, is_difficult2), axis=0)
- result['difficult'] = is_difficult
- return result
- class RandomDistort(Transform):
- """
- Random color distortion.
- Args:
- brightness_range (float, optional): Range of brightness distortion.
- Defaults to .5.
- brightness_prob (float, optional): Probability of brightness distortion.
- Defaults to .5.
- contrast_range (float, optional): Range of contrast distortion.
- Defaults to .5.
- contrast_prob (float, optional): Probability of contrast distortion.
- Defaults to .5.
- saturation_range (float, optional): Range of saturation distortion.
- Defaults to .5.
- saturation_prob (float, optional): Probability of saturation distortion.
- Defaults to .5.
- hue_range (float, optional): Range of hue distortion. Defaults to .5.
- hue_prob (float, optional): Probability of hue distortion. Defaults to .5.
- random_apply (bool, optional): Apply the transformation in random (yolo) or
- fixed (SSD) order. Defaults to True.
- count (int, optional): Number of distortions to apply. Defaults to 4.
- shuffle_channel (bool, optional): Whether to swap channels randomly.
- Defaults to False.
- """
- def __init__(self,
- brightness_range=0.5,
- brightness_prob=0.5,
- contrast_range=0.5,
- contrast_prob=0.5,
- saturation_range=0.5,
- saturation_prob=0.5,
- hue_range=18,
- hue_prob=0.5,
- random_apply=True,
- count=4,
- shuffle_channel=False):
- super(RandomDistort, self).__init__()
- self.brightness_range = [1 - brightness_range, 1 + brightness_range]
- self.brightness_prob = brightness_prob
- self.contrast_range = [1 - contrast_range, 1 + contrast_range]
- self.contrast_prob = contrast_prob
- self.saturation_range = [1 - saturation_range, 1 + saturation_range]
- self.saturation_prob = saturation_prob
- self.hue_range = [1 - hue_range, 1 + hue_range]
- self.hue_prob = hue_prob
- self.random_apply = random_apply
- self.count = count
- self.shuffle_channel = shuffle_channel
- def apply_hue(self, image):
- low, high = self.hue_range
- if np.random.uniform(0., 1.) < self.hue_prob:
- return image
-
- delta = np.random.uniform(low, high)
- u = np.cos(delta * np.pi)
- w = np.sin(delta * np.pi)
- bt = np.array([[1.0, 0.0, 0.0], [0.0, u, -w], [0.0, w, u]])
- tyiq = np.array([[0.299, 0.587, 0.114], [0.596, -0.274, -0.321],
- [0.211, -0.523, 0.311]])
- ityiq = np.array([[1.0, 0.956, 0.621], [1.0, -0.272, -0.647],
- [1.0, -1.107, 1.705]])
- t = np.dot(np.dot(ityiq, bt), tyiq).T
- res_list = []
- channel = image.shape[2]
- for i in range(channel // 3):
- sub_img = image[:, :, 3 * i:3 * (i + 1)]
- sub_img = sub_img.astype(np.float32)
- sub_img = np.dot(image, t)
- res_list.append(sub_img)
- if channel % 3 != 0:
- i = channel % 3
- res_list.append(image[:, :, -i:])
- return np.concatenate(res_list, axis=2)
- def apply_saturation(self, image):
- low, high = self.saturation_range
- delta = np.random.uniform(low, high)
- if np.random.uniform(0., 1.) < self.saturation_prob:
- return image
- res_list = []
- channel = image.shape[2]
- for i in range(channel // 3):
- sub_img = image[:, :, 3 * i:3 * (i + 1)]
- sub_img = sub_img.astype(np.float32)
-
- gray = sub_img * np.array(
- [[[0.299, 0.587, 0.114]]], dtype=np.float32)
- gray = gray.sum(axis=2, keepdims=True)
- gray *= (1.0 - delta)
- sub_img *= delta
- sub_img += gray
- res_list.append(sub_img)
- if channel % 3 != 0:
- i = channel % 3
- res_list.append(image[:, :, -i:])
- return np.concatenate(res_list, axis=2)
- def apply_contrast(self, image):
- low, high = self.contrast_range
- if np.random.uniform(0., 1.) < self.contrast_prob:
- return image
- delta = np.random.uniform(low, high)
- image = image.astype(np.float32)
- image *= delta
- return image
- def apply_brightness(self, image):
- low, high = self.brightness_range
- if np.random.uniform(0., 1.) < self.brightness_prob:
- return image
- delta = np.random.uniform(low, high)
- image = image.astype(np.float32)
- image += delta
- return image
- def apply(self, sample):
- if self.random_apply:
- functions = [
- self.apply_brightness, self.apply_contrast,
- self.apply_saturation, self.apply_hue
- ]
- distortions = np.random.permutation(functions)[:self.count]
- for func in distortions:
- sample['image'] = func(sample['image'])
- if 'image2' in sample:
- sample['image2'] = func(sample['image2'])
- return sample
- sample['image'] = self.apply_brightness(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_brightness(sample['image2'])
- mode = np.random.randint(0, 2)
- if mode:
- sample['image'] = self.apply_contrast(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_contrast(sample['image2'])
- sample['image'] = self.apply_saturation(sample['image'])
- sample['image'] = self.apply_hue(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_saturation(sample['image2'])
- sample['image2'] = self.apply_hue(sample['image2'])
- if not mode:
- sample['image'] = self.apply_contrast(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_contrast(sample['image2'])
- if self.shuffle_channel:
- if np.random.randint(0, 2):
- sample['image'] = sample['image'][..., np.random.permutation(3)]
- if 'image2' in sample:
- sample['image2'] = sample['image2'][
- ..., np.random.permutation(3)]
- return sample
- class RandomBlur(Transform):
- """
- Randomly blur input image(s).
- Args:
- prob (float): Probability of blurring.
- """
- def __init__(self, prob=0.1):
- super(RandomBlur, self).__init__()
- self.prob = prob
- def apply_im(self, image, radius):
- image = cv2.GaussianBlur(image, (radius, radius), 0, 0)
- return image
- def apply(self, sample):
- if self.prob <= 0:
- n = 0
- elif self.prob >= 1:
- n = 1
- else:
- n = int(1.0 / self.prob)
- if n > 0:
- if np.random.randint(0, n) == 0:
- radius = np.random.randint(3, 10)
- if radius % 2 != 1:
- radius = radius + 1
- if radius > 9:
- radius = 9
- sample['image'] = self.apply_im(sample['image'], radius)
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'], radius)
- return sample
- class Dehaze(Transform):
- """
- Dehaze input image(s).
- Args:
- gamma (bool, optional): Use gamma correction or not. Defaults to False.
- """
- def __init__(self, gamma=False):
- super(Dehaze, self).__init__()
- self.gamma = gamma
- def apply_im(self, image):
- image = dehaze(image, self.gamma)
- return image
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- return sample
- class ReduceDim(Transform):
- """
- Use PCA to reduce the dimension of input image(s).
- Args:
- joblib_path (str): Path of *.joblib file of PCA.
- apply_to_tar (bool, optional): Whether to apply transformation to the target
- image. Defaults to True.
- """
- def __init__(self, joblib_path, apply_to_tar=True):
- super(ReduceDim, self).__init__()
- ext = joblib_path.split(".")[-1]
- if ext != "joblib":
- raise ValueError("`joblib_path` must be *.joblib, not *.{}.".format(
- ext))
- self.pca = load(joblib_path)
- self.apply_to_tar = apply_to_tar
- def apply_im(self, image):
- H, W, C = image.shape
- n_im = np.reshape(image, (-1, C))
- im_pca = self.pca.transform(n_im)
- result = np.reshape(im_pca, (H, W, -1))
- return result
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'target' in sample and self.apply_to_tar:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class SelectBand(Transform):
- """
- Select a set of bands of input image(s).
- Args:
- band_list (list, optional): Bands to select (band index starts from 1).
- Defaults to [1, 2, 3].
- apply_to_tar (bool, optional): Whether to apply transformation to the target
- image. Defaults to True.
- """
- def __init__(self, band_list=[1, 2, 3], apply_to_tar=True):
- super(SelectBand, self).__init__()
- self.band_list = band_list
- self.apply_to_tar = apply_to_tar
- def apply_im(self, image):
- image = select_bands(image, self.band_list)
- return image
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- if 'target' in sample and self.apply_to_tar:
- sample['target'] = self.apply_im(sample['target'])
- return sample
- class _PadBox(Transform):
- def __init__(self, num_max_boxes=50):
- """
- Pad zeros to bboxes if number of bboxes is less than `num_max_boxes`.
- Args:
- num_max_boxes (int, optional): Max number of bboxes. Defaults to 50.
- """
- self.num_max_boxes = num_max_boxes
- super(_PadBox, self).__init__()
- def apply(self, sample):
- gt_num = min(self.num_max_boxes, len(sample['gt_bbox']))
- num_max = self.num_max_boxes
- pad_bbox = np.zeros((num_max, 4), dtype=np.float32)
- if gt_num > 0:
- pad_bbox[:gt_num, :] = sample['gt_bbox'][:gt_num, :]
- sample['gt_bbox'] = pad_bbox
- if 'gt_class' in sample:
- pad_class = np.zeros((num_max, ), dtype=np.int32)
- if gt_num > 0:
- pad_class[:gt_num] = sample['gt_class'][:gt_num, 0]
- sample['gt_class'] = pad_class
- if 'gt_score' in sample:
- pad_score = np.zeros((num_max, ), dtype=np.float32)
- if gt_num > 0:
- pad_score[:gt_num] = sample['gt_score'][:gt_num, 0]
- sample['gt_score'] = pad_score
-
-
-
- if 'difficult' in sample:
- pad_diff = np.zeros((num_max, ), dtype=np.int32)
- if gt_num > 0:
- pad_diff[:gt_num] = sample['difficult'][:gt_num, 0]
- sample['difficult'] = pad_diff
- if 'is_crowd' in sample:
- pad_crowd = np.zeros((num_max, ), dtype=np.int32)
- if gt_num > 0:
- pad_crowd[:gt_num] = sample['is_crowd'][:gt_num, 0]
- sample['is_crowd'] = pad_crowd
- return sample
- class _NormalizeBox(Transform):
- def __init__(self):
- super(_NormalizeBox, self).__init__()
- def apply(self, sample):
- height, width = sample['image'].shape[:2]
- for i in range(sample['gt_bbox'].shape[0]):
- sample['gt_bbox'][i][0] = sample['gt_bbox'][i][0] / width
- sample['gt_bbox'][i][1] = sample['gt_bbox'][i][1] / height
- sample['gt_bbox'][i][2] = sample['gt_bbox'][i][2] / width
- sample['gt_bbox'][i][3] = sample['gt_bbox'][i][3] / height
- return sample
- class _BboxXYXY2XYWH(Transform):
- """
- Convert bbox XYXY format to XYWH format.
- """
- def __init__(self):
- super(_BboxXYXY2XYWH, self).__init__()
- def apply(self, sample):
- bbox = sample['gt_bbox']
- bbox[:, 2:4] = bbox[:, 2:4] - bbox[:, :2]
- bbox[:, :2] = bbox[:, :2] + bbox[:, 2:4] / 2.
- sample['gt_bbox'] = bbox
- return sample
- class _Permute(Transform):
- def __init__(self):
- super(_Permute, self).__init__()
- def apply(self, sample):
- sample['image'] = permute(sample['image'], False)
- if 'image2' in sample:
- sample['image2'] = permute(sample['image2'], False)
- if 'target' in sample:
- sample['target'] = permute(sample['target'], False)
- return sample
- class RandomSwap(Transform):
- """
- Randomly swap multi-temporal images.
- Args:
- prob (float, optional): Probability of swapping the input images.
- Default: 0.2.
- """
- def __init__(self, prob=0.2):
- super(RandomSwap, self).__init__()
- self.prob = prob
- def apply(self, sample):
- if 'image2' not in sample:
- raise ValueError("'image2' is not found in the sample.")
- if random.random() < self.prob:
- sample['image'], sample['image2'] = sample['image2'], sample[
- 'image']
- return sample
- class ReloadMask(Transform):
- def apply(self, sample):
- sample['mask'] = decode_seg_mask(sample['mask_ori'])
- if 'aux_masks' in sample:
- sample['aux_masks'] = list(
- map(decode_seg_mask, sample['aux_masks_ori']))
- return sample
- class AppendIndex(Transform):
- """
- Append remote sensing index to input image(s).
- Args:
- index_type (str): Type of remote sensinng index. See supported
- index types in
- https://github.com/PaddlePaddle/PaddleRS/tree/develop/paddlers/transforms/indices.py .
- band_indices (dict, optional): Mapping of band names to band indices
- (starting from 1). See band names in
- https://github.com/PaddlePaddle/PaddleRS/tree/develop/paddlers/transforms/indices.py .
- Default: None.
- satellite (str, optional): Type of satellite. If set,
- band indices will be automatically determined accordingly. See supported satellites in
- https://github.com/PaddlePaddle/PaddleRS/tree/develop/paddlers/transforms/satellites.py .
- Default: None.
- """
- def __init__(self, index_type, band_indices=None, satellite=None, **kwargs):
- super(AppendIndex, self).__init__()
- cls = getattr(indices, index_type)
- if satellite is not None:
- satellite_bands = getattr(satellites, satellite)
- self._compute_index = cls(satellite_bands, **kwargs)
- else:
- if band_indices is None:
- raise ValueError(
- "At least one of `band_indices` and `satellite` must not be None."
- )
- else:
- self._compute_index = cls(band_indices, **kwargs)
- def apply_im(self, image):
- index = self._compute_index(image)
- index = index[..., None].astype('float32')
- return np.concatenate([image, index], axis=-1)
- def apply(self, sample):
- sample['image'] = self.apply_im(sample['image'])
- if 'image2' in sample:
- sample['image2'] = self.apply_im(sample['image2'])
- return sample
- class MatchRadiance(Transform):
- """
- Perform relative radiometric correction between bi-temporal images.
- Args:
- method (str, optional): Method used to match the radiance of the
- bi-temporal images. Choices are {'hist', 'lsr'}. 'hist' stands
- for histogram matching and 'lsr' stands for least-squares
- regression. Default: 'hist'.
- """
- def __init__(self, method='hist'):
- super(MatchRadiance, self).__init__()
- if method == 'hist':
- self._match_func = match_histograms
- elif method == 'lsr':
- self._match_func = match_by_regression
- else:
- raise ValueError(
- "{} is not a supported radiometric correction method.".format(
- method))
- self.method = method
- def apply(self, sample):
- if 'image2' not in sample:
- raise ValueError("'image2' is not found in the sample.")
- sample['image2'] = self._match_func(sample['image2'], sample['image'])
- return sample
- class Arrange(Transform):
- def __init__(self, mode):
- super().__init__()
- if mode not in ['train', 'eval', 'test', 'quant']:
- raise ValueError(
- "`mode` should be defined as one of ['train', 'eval', 'test', 'quant']!"
- )
- self.mode = mode
- class ArrangeSegmenter(Arrange):
- def apply(self, sample):
- if 'mask' in sample:
- mask = sample['mask']
- mask = mask.astype('int64')
- image = permute(sample['image'], False)
- if self.mode == 'train':
- return image, mask
- if self.mode == 'eval':
- return image, mask
- if self.mode == 'test':
- return image,
- class ArrangeChangeDetector(Arrange):
- def apply(self, sample):
- if 'mask' in sample:
- mask = sample['mask']
- mask = mask.astype('int64')
- image_t1 = permute(sample['image'], False)
- image_t2 = permute(sample['image2'], False)
- if self.mode == 'train':
- masks = [mask]
- if 'aux_masks' in sample:
- masks.extend(
- map(methodcaller('astype', 'int64'), sample['aux_masks']))
- return (
- image_t1,
- image_t2, ) + tuple(masks)
- if self.mode == 'eval':
- return image_t1, image_t2, mask
- if self.mode == 'test':
- return image_t1, image_t2,
- class ArrangeClassifier(Arrange):
- def apply(self, sample):
- image = permute(sample['image'], False)
- if self.mode in ['train', 'eval']:
- return image, sample['label']
- else:
- return image
- class ArrangeDetector(Arrange):
- def apply(self, sample):
- if self.mode == 'eval' and 'gt_poly' in sample:
- del sample['gt_poly']
- return sample
- class ArrangeRestorer(Arrange):
- def apply(self, sample):
- if 'target' in sample:
- target = permute(sample['target'], False)
- image = permute(sample['image'], False)
- if self.mode == 'train':
- return image, target
- if self.mode == 'eval':
- return image, target
- if self.mode == 'test':
- return image,
|