(self, height, width, border_width)
| 8 | |
| 9 | |
| 10 | def mask(self, height, width, border_width): |
| 11 | # Create a mask with shape (height, width). |
| 12 | # The centre area is filled with 1, and the border line is filled with values in range (0, 1]. |
| 13 | x = torch.arange(height).repeat(width, 1).T |
| 14 | y = torch.arange(width).repeat(height, 1) |
| 15 | mask = torch.stack([x + 1, height - x, y + 1, width - y]).min(dim=0).values |
| 16 | mask = (mask / border_width).clip(0, 1) |
| 17 | return mask |
| 18 | |
| 19 | |
| 20 | def tile(self, model_input, tile_size, tile_stride, tile_device, tile_dtype): |