Spaces:
Running
on
T4
Running
on
T4
| """ | |
| Copyright (C) 2018 NVIDIA Corporation. All rights reserved. | |
| Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). | |
| """ | |
| src = ''' | |
| #include "/usr/local/cuda/include/math_functions.h" | |
| #define TB 256 | |
| #define EPS 1e-7 | |
| __device__ bool InverseMat4x4(double m_in[4][4], double inv_out[4][4]) { | |
| double m[16], inv[16]; | |
| for (int i = 0; i < 4; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| m[i * 4 + j] = m_in[i][j]; | |
| } | |
| } | |
| inv[0] = m[5] * m[10] * m[15] - | |
| m[5] * m[11] * m[14] - | |
| m[9] * m[6] * m[15] + | |
| m[9] * m[7] * m[14] + | |
| m[13] * m[6] * m[11] - | |
| m[13] * m[7] * m[10]; | |
| inv[4] = -m[4] * m[10] * m[15] + | |
| m[4] * m[11] * m[14] + | |
| m[8] * m[6] * m[15] - | |
| m[8] * m[7] * m[14] - | |
| m[12] * m[6] * m[11] + | |
| m[12] * m[7] * m[10]; | |
| inv[8] = m[4] * m[9] * m[15] - | |
| m[4] * m[11] * m[13] - | |
| m[8] * m[5] * m[15] + | |
| m[8] * m[7] * m[13] + | |
| m[12] * m[5] * m[11] - | |
| m[12] * m[7] * m[9]; | |
| inv[12] = -m[4] * m[9] * m[14] + | |
| m[4] * m[10] * m[13] + | |
| m[8] * m[5] * m[14] - | |
| m[8] * m[6] * m[13] - | |
| m[12] * m[5] * m[10] + | |
| m[12] * m[6] * m[9]; | |
| inv[1] = -m[1] * m[10] * m[15] + | |
| m[1] * m[11] * m[14] + | |
| m[9] * m[2] * m[15] - | |
| m[9] * m[3] * m[14] - | |
| m[13] * m[2] * m[11] + | |
| m[13] * m[3] * m[10]; | |
| inv[5] = m[0] * m[10] * m[15] - | |
| m[0] * m[11] * m[14] - | |
| m[8] * m[2] * m[15] + | |
| m[8] * m[3] * m[14] + | |
| m[12] * m[2] * m[11] - | |
| m[12] * m[3] * m[10]; | |
| inv[9] = -m[0] * m[9] * m[15] + | |
| m[0] * m[11] * m[13] + | |
| m[8] * m[1] * m[15] - | |
| m[8] * m[3] * m[13] - | |
| m[12] * m[1] * m[11] + | |
| m[12] * m[3] * m[9]; | |
| inv[13] = m[0] * m[9] * m[14] - | |
| m[0] * m[10] * m[13] - | |
| m[8] * m[1] * m[14] + | |
| m[8] * m[2] * m[13] + | |
| m[12] * m[1] * m[10] - | |
| m[12] * m[2] * m[9]; | |
| inv[2] = m[1] * m[6] * m[15] - | |
| m[1] * m[7] * m[14] - | |
| m[5] * m[2] * m[15] + | |
| m[5] * m[3] * m[14] + | |
| m[13] * m[2] * m[7] - | |
| m[13] * m[3] * m[6]; | |
| inv[6] = -m[0] * m[6] * m[15] + | |
| m[0] * m[7] * m[14] + | |
| m[4] * m[2] * m[15] - | |
| m[4] * m[3] * m[14] - | |
| m[12] * m[2] * m[7] + | |
| m[12] * m[3] * m[6]; | |
| inv[10] = m[0] * m[5] * m[15] - | |
| m[0] * m[7] * m[13] - | |
| m[4] * m[1] * m[15] + | |
| m[4] * m[3] * m[13] + | |
| m[12] * m[1] * m[7] - | |
| m[12] * m[3] * m[5]; | |
| inv[14] = -m[0] * m[5] * m[14] + | |
| m[0] * m[6] * m[13] + | |
| m[4] * m[1] * m[14] - | |
| m[4] * m[2] * m[13] - | |
| m[12] * m[1] * m[6] + | |
| m[12] * m[2] * m[5]; | |
| inv[3] = -m[1] * m[6] * m[11] + | |
| m[1] * m[7] * m[10] + | |
| m[5] * m[2] * m[11] - | |
| m[5] * m[3] * m[10] - | |
| m[9] * m[2] * m[7] + | |
| m[9] * m[3] * m[6]; | |
| inv[7] = m[0] * m[6] * m[11] - | |
| m[0] * m[7] * m[10] - | |
| m[4] * m[2] * m[11] + | |
| m[4] * m[3] * m[10] + | |
| m[8] * m[2] * m[7] - | |
| m[8] * m[3] * m[6]; | |
| inv[11] = -m[0] * m[5] * m[11] + | |
| m[0] * m[7] * m[9] + | |
| m[4] * m[1] * m[11] - | |
| m[4] * m[3] * m[9] - | |
| m[8] * m[1] * m[7] + | |
| m[8] * m[3] * m[5]; | |
| inv[15] = m[0] * m[5] * m[10] - | |
| m[0] * m[6] * m[9] - | |
| m[4] * m[1] * m[10] + | |
| m[4] * m[2] * m[9] + | |
| m[8] * m[1] * m[6] - | |
| m[8] * m[2] * m[5]; | |
| double det = m[0] * inv[0] + m[1] * inv[4] + m[2] * inv[8] + m[3] * inv[12]; | |
| if (abs(det) < 1e-9) { | |
| return false; | |
| } | |
| det = 1.0 / det; | |
| for (int i = 0; i < 4; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| inv_out[i][j] = inv[i * 4 + j] * det; | |
| } | |
| } | |
| return true; | |
| } | |
| extern "C" | |
| __global__ void best_local_affine_kernel( | |
| float *output, float *input, float *affine_model, | |
| int h, int w, float epsilon, int kernel_radius | |
| ) | |
| { | |
| int size = h * w; | |
| int id = blockIdx.x * blockDim.x + threadIdx.x; | |
| if (id < size) { | |
| int x = id % w, y = id / w; | |
| double Mt_M[4][4] = {}; // 4x4 | |
| double invMt_M[4][4] = {}; | |
| double Mt_S[3][4] = {}; // RGB -> 1x4 | |
| double A[3][4] = {}; | |
| for (int i = 0; i < 4; i++) | |
| for (int j = 0; j < 4; j++) { | |
| Mt_M[i][j] = 0, invMt_M[i][j] = 0; | |
| if (i != 3) { | |
| Mt_S[i][j] = 0, A[i][j] = 0; | |
| if (i == j) | |
| Mt_M[i][j] = 1e-3; | |
| } | |
| } | |
| for (int dy = -kernel_radius; dy <= kernel_radius; dy++) { | |
| for (int dx = -kernel_radius; dx <= kernel_radius; dx++) { | |
| int xx = x + dx, yy = y + dy; | |
| int id2 = yy * w + xx; | |
| if (0 <= xx && xx < w && 0 <= yy && yy < h) { | |
| Mt_M[0][0] += input[id2 + 2*size] * input[id2 + 2*size]; | |
| Mt_M[0][1] += input[id2 + 2*size] * input[id2 + size]; | |
| Mt_M[0][2] += input[id2 + 2*size] * input[id2]; | |
| Mt_M[0][3] += input[id2 + 2*size]; | |
| Mt_M[1][0] += input[id2 + size] * input[id2 + 2*size]; | |
| Mt_M[1][1] += input[id2 + size] * input[id2 + size]; | |
| Mt_M[1][2] += input[id2 + size] * input[id2]; | |
| Mt_M[1][3] += input[id2 + size]; | |
| Mt_M[2][0] += input[id2] * input[id2 + 2*size]; | |
| Mt_M[2][1] += input[id2] * input[id2 + size]; | |
| Mt_M[2][2] += input[id2] * input[id2]; | |
| Mt_M[2][3] += input[id2]; | |
| Mt_M[3][0] += input[id2 + 2*size]; | |
| Mt_M[3][1] += input[id2 + size]; | |
| Mt_M[3][2] += input[id2]; | |
| Mt_M[3][3] += 1; | |
| Mt_S[0][0] += input[id2 + 2*size] * output[id2 + 2*size]; | |
| Mt_S[0][1] += input[id2 + size] * output[id2 + 2*size]; | |
| Mt_S[0][2] += input[id2] * output[id2 + 2*size]; | |
| Mt_S[0][3] += output[id2 + 2*size]; | |
| Mt_S[1][0] += input[id2 + 2*size] * output[id2 + size]; | |
| Mt_S[1][1] += input[id2 + size] * output[id2 + size]; | |
| Mt_S[1][2] += input[id2] * output[id2 + size]; | |
| Mt_S[1][3] += output[id2 + size]; | |
| Mt_S[2][0] += input[id2 + 2*size] * output[id2]; | |
| Mt_S[2][1] += input[id2 + size] * output[id2]; | |
| Mt_S[2][2] += input[id2] * output[id2]; | |
| Mt_S[2][3] += output[id2]; | |
| } | |
| } | |
| } | |
| bool success = InverseMat4x4(Mt_M, invMt_M); | |
| for (int i = 0; i < 3; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| for (int k = 0; k < 4; k++) { | |
| A[i][j] += invMt_M[j][k] * Mt_S[i][k]; | |
| } | |
| } | |
| } | |
| for (int i = 0; i < 3; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| int affine_id = i * 4 + j; | |
| affine_model[12 * id + affine_id] = A[i][j]; | |
| } | |
| } | |
| } | |
| return ; | |
| } | |
| extern "C" | |
| __global__ void bilateral_smooth_kernel( | |
| float *affine_model, float *filtered_affine_model, float *guide, | |
| int h, int w, int kernel_radius, float sigma1, float sigma2 | |
| ) | |
| { | |
| int id = blockIdx.x * blockDim.x + threadIdx.x; | |
| int size = h * w; | |
| if (id < size) { | |
| int x = id % w; | |
| int y = id / w; | |
| double sum_affine[12] = {}; | |
| double sum_weight = 0; | |
| for (int dx = -kernel_radius; dx <= kernel_radius; dx++) { | |
| for (int dy = -kernel_radius; dy <= kernel_radius; dy++) { | |
| int yy = y + dy, xx = x + dx; | |
| int id2 = yy * w + xx; | |
| if (0 <= xx && xx < w && 0 <= yy && yy < h) { | |
| float color_diff1 = guide[yy*w + xx] - guide[y*w + x]; | |
| float color_diff2 = guide[yy*w + xx + size] - guide[y*w + x + size]; | |
| float color_diff3 = guide[yy*w + xx + 2*size] - guide[y*w + x + 2*size]; | |
| float color_diff_sqr = | |
| (color_diff1*color_diff1 + color_diff2*color_diff2 + color_diff3*color_diff3) / 3; | |
| float v1 = exp(-(dx * dx + dy * dy) / (2 * sigma1 * sigma1)); | |
| float v2 = exp(-(color_diff_sqr) / (2 * sigma2 * sigma2)); | |
| float weight = v1 * v2; | |
| for (int i = 0; i < 3; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| int affine_id = i * 4 + j; | |
| sum_affine[affine_id] += weight * affine_model[id2*12 + affine_id]; | |
| } | |
| } | |
| sum_weight += weight; | |
| } | |
| } | |
| } | |
| for (int i = 0; i < 3; i++) { | |
| for (int j = 0; j < 4; j++) { | |
| int affine_id = i * 4 + j; | |
| filtered_affine_model[id*12 + affine_id] = sum_affine[affine_id] / sum_weight; | |
| } | |
| } | |
| } | |
| return ; | |
| } | |
| extern "C" | |
| __global__ void reconstruction_best_kernel( | |
| float *input, float *filtered_affine_model, float *filtered_best_output, | |
| int h, int w | |
| ) | |
| { | |
| int id = blockIdx.x * blockDim.x + threadIdx.x; | |
| int size = h * w; | |
| if (id < size) { | |
| double out1 = | |
| input[id + 2*size] * filtered_affine_model[id*12 + 0] + // A[0][0] + | |
| input[id + size] * filtered_affine_model[id*12 + 1] + // A[0][1] + | |
| input[id] * filtered_affine_model[id*12 + 2] + // A[0][2] + | |
| filtered_affine_model[id*12 + 3]; //A[0][3]; | |
| double out2 = | |
| input[id + 2*size] * filtered_affine_model[id*12 + 4] + //A[1][0] + | |
| input[id + size] * filtered_affine_model[id*12 + 5] + //A[1][1] + | |
| input[id] * filtered_affine_model[id*12 + 6] + //A[1][2] + | |
| filtered_affine_model[id*12 + 7]; //A[1][3]; | |
| double out3 = | |
| input[id + 2*size] * filtered_affine_model[id*12 + 8] + //A[2][0] + | |
| input[id + size] * filtered_affine_model[id*12 + 9] + //A[2][1] + | |
| input[id] * filtered_affine_model[id*12 + 10] + //A[2][2] + | |
| filtered_affine_model[id*12 + 11]; // A[2][3]; | |
| filtered_best_output[id] = out1; | |
| filtered_best_output[id + size] = out2; | |
| filtered_best_output[id + 2*size] = out3; | |
| } | |
| return ; | |
| } | |
| ''' | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from cupy.cuda import function | |
| from pynvrtc.compiler import Program | |
| from collections import namedtuple | |
| def smooth_local_affine(output_cpu, input_cpu, epsilon, patch, h, w, f_r, f_e): | |
| # program = Program(src.encode('utf-8'), 'best_local_affine_kernel.cu'.encode('utf-8')) | |
| # ptx = program.compile(['-I/usr/local/cuda/include'.encode('utf-8')]) | |
| program = Program(src, 'best_local_affine_kernel.cu') | |
| ptx = program.compile(['-I/usr/local/cuda/include']) | |
| m = function.Module() | |
| m.load(bytes(ptx.encode())) | |
| _reconstruction_best_kernel = m.get_function('reconstruction_best_kernel') | |
| _bilateral_smooth_kernel = m.get_function('bilateral_smooth_kernel') | |
| _best_local_affine_kernel = m.get_function('best_local_affine_kernel') | |
| Stream = namedtuple('Stream', ['ptr']) | |
| s = Stream(ptr=torch.cuda.current_stream().cuda_stream) | |
| filter_radius = f_r | |
| sigma1 = filter_radius / 3 | |
| sigma2 = f_e | |
| radius = (patch - 1) / 2 | |
| filtered_best_output = torch.zeros(np.shape(input_cpu)).cuda() | |
| affine_model = torch.zeros((h * w, 12)).cuda() | |
| filtered_affine_model =torch.zeros((h * w, 12)).cuda() | |
| input_ = torch.from_numpy(input_cpu).cuda() | |
| output_ = torch.from_numpy(output_cpu).cuda() | |
| _best_local_affine_kernel( | |
| grid=(int((h * w) / 256 + 1), 1), | |
| block=(256, 1, 1), | |
| args=[output_.data_ptr(), input_.data_ptr(), affine_model.data_ptr(), | |
| np.int32(h), np.int32(w), np.float32(epsilon), np.int32(radius)], stream=s | |
| ) | |
| _bilateral_smooth_kernel( | |
| grid=(int((h * w) / 256 + 1), 1), | |
| block=(256, 1, 1), | |
| args=[affine_model.data_ptr(), filtered_affine_model.data_ptr(), input_.data_ptr(), np.int32(h), np.int32(w), np.int32(f_r), np.float32(sigma1), np.float32(sigma2)], stream=s | |
| ) | |
| _reconstruction_best_kernel( | |
| grid=(int((h * w) / 256 + 1), 1), | |
| block=(256, 1, 1), | |
| args=[input_.data_ptr(), filtered_affine_model.data_ptr(), filtered_best_output.data_ptr(), | |
| np.int32(h), np.int32(w)], stream=s | |
| ) | |
| numpy_filtered_best_output = filtered_best_output.cpu().numpy() | |
| return numpy_filtered_best_output | |
| def smooth_filter(initImg, contentImg, f_radius=15,f_edge=1e-1): | |
| ''' | |
| :param initImg: intermediate output. Either image path or PIL Image | |
| :param contentImg: content image output. Either path or PIL Image | |
| :return: stylized output image. PIL Image | |
| ''' | |
| if type(initImg) == str: | |
| initImg = Image.open(initImg).convert("RGB") | |
| best_image_bgr = np.array(initImg, dtype=np.float32) | |
| bW, bH, bC = best_image_bgr.shape | |
| best_image_bgr = best_image_bgr[:, :, ::-1] | |
| best_image_bgr = best_image_bgr.transpose((2, 0, 1)) | |
| if type(contentImg) == str: | |
| contentImg = Image.open(contentImg).convert("RGB") | |
| content_input = contentImg.resize((bH,bW)) | |
| content_input = np.array(content_input, dtype=np.float32) | |
| content_input = content_input[:, :, ::-1] | |
| content_input = content_input.transpose((2, 0, 1)) | |
| input_ = np.ascontiguousarray(content_input, dtype=np.float32) / 255. | |
| _, H, W = np.shape(input_) | |
| output_ = np.ascontiguousarray(best_image_bgr, dtype=np.float32) / 255. | |
| best_ = smooth_local_affine(output_, input_, 1e-7, 3, H, W, f_radius, f_edge) | |
| best_ = best_.transpose(1, 2, 0) | |
| result = Image.fromarray(np.uint8(np.clip(best_ * 255., 0, 255.))) | |
| return result | |