| 568 | |
| 569 | |
| 570 | class MaskTransform(ImageTransform): |
| 571 | |
| 572 | def __init__(self, mask_pool_size=1): |
| 573 | assert isinstance(mask_pool_size, int) |
| 574 | self.mask_pool_size = mask_pool_size # Use to expand masks |
| 575 | |
| 576 | def mask_to_tensor(self, img): |
| 577 | mask = TF.to_tensor(img) |
| 578 | if self.mask_pool_size > 1: |
| 579 | mask = reduce(mask, 'c (h1 h2) (w1 w2) -> c h1 w1', 'min', h2=self.mask_pool_size, w2=self.mask_pool_size) |
| 580 | mask = repeat(mask, 'c h1 w1 -> c (h1 h2) (w1 w2)', h2=self.mask_pool_size, w2=self.mask_pool_size) |
| 581 | return (mask == 1.0) |
| 582 | |
| 583 | def load(self, path): |
| 584 | sample = self.pil_loader(path) |
| 585 | return sample |
| 586 | |
| 587 | def preprocess(self, sample): |
| 588 | return sample |
| 589 | |
| 590 | def image_augment(self, img, crop_coords: Tuple, flip: bool, orig_size: Tuple, target_size: Tuple, |
| 591 | rand_aug_idx: Optional[int], resample_mode: str = None): |
| 592 | # Override resampling mode to 'nearest' for masks |
| 593 | img = self.image_crop_and_resize(img, crop_coords, target_size, resample_mode='nearest') |
| 594 | img = self.image_hflip(img, flip) |
| 595 | return img |
| 596 | |
| 597 | def postprocess(self, sample): |
| 598 | sample = self.mask_to_tensor(sample) |
| 599 | return sample |
| 600 | |
| 601 | |
| 602 | class TokTransform(AbstractTransform): |
no outgoing calls
no test coverage detected