| 22 | |
| 23 | |
| 24 | torch::Tensor push_pull_inpaint( |
| 25 | const torch::Tensor& mask, |
| 26 | const torch::Tensor& data) |
| 27 | { |
| 28 | //check input |
| 29 | CHECK_CUDA(mask); |
| 30 | CHECK_CONTIGUOUS(mask); |
| 31 | CHECK_DIM(mask, 3); |
| 32 | int64_t B = mask.size(0); |
| 33 | int64_t H = mask.size(1); |
| 34 | int64_t W = mask.size(2); |
| 35 | |
| 36 | CHECK_CUDA(data); |
| 37 | CHECK_CONTIGUOUS(data); |
| 38 | CHECK_DIM(data, 4); |
| 39 | CHECK_SIZE(data, 0, B); |
| 40 | int64_t C = data.size(1); |
| 41 | CHECK_SIZE(data, 2, H); |
| 42 | CHECK_SIZE(data, 3, W); |
| 43 | TORCH_CHECK(C < MAX_CHANNELS, "Inpainting::fastInpaint only supports up to 128 channels, but got " + std::to_string(C)); |
| 44 | |
| 45 | //inpaint recursivly |
| 46 | return std::get<1>(push_pull_inpaint_recursion_cuda(mask, data)); |
| 47 | } |
| 48 | |
| 49 | |
| 50 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |