| 18 | |
| 19 | |
| 20 | class RandomCrop(object): |
| 21 | def __init__(self, size, padding=0): |
| 22 | if isinstance(size, numbers.Number): |
| 23 | self.size = (int(size), int(size)) |
| 24 | else: |
| 25 | self.size = size |
| 26 | self.padding = padding |
| 27 | |
| 28 | def __call__(self, img, mask): |
| 29 | if self.padding > 0: |
| 30 | img = ImageOps.expand(img, border=self.padding, fill=0) |
| 31 | mask = ImageOps.expand(mask, border=self.padding, fill=0) |
| 32 | |
| 33 | assert img.size == mask.size |
| 34 | w, h = img.size |
| 35 | th, tw = self.size |
| 36 | if w == tw and h == th: |
| 37 | return img, mask |
| 38 | if w < tw or h < th: |
| 39 | return img.resize((tw, th), Image.BILINEAR), mask.resize((tw, th), Image.NEAREST) |
| 40 | |
| 41 | x1 = random.randint(0, w - tw) |
| 42 | y1 = random.randint(0, h - th) |
| 43 | return img.crop((x1, y1, x1 + tw, y1 + th)), mask.crop((x1, y1, x1 + tw, y1 + th)) |
| 44 | |
| 45 | |
| 46 | class CenterCrop(object): |