Correlationn operator in NCHW layout. Parameters ---------- data1_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] data2_np : numpy.ndarray 4-D with shape [batch, in_channel, in_height, in_width] kernel_size: int Kernel size fo
(
data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply
)
| 22 | |
| 23 | |
| 24 | def correlation_nchw_python( |
| 25 | data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply |
| 26 | ): |
| 27 | """Correlationn operator in NCHW layout. |
| 28 | |
| 29 | Parameters |
| 30 | ---------- |
| 31 | data1_np : numpy.ndarray |
| 32 | 4-D with shape [batch, in_channel, in_height, in_width] |
| 33 | |
| 34 | data2_np : numpy.ndarray |
| 35 | 4-D with shape [batch, in_channel, in_height, in_width] |
| 36 | |
| 37 | kernel_size: int |
| 38 | Kernel size for correlation, must be an odd number |
| 39 | |
| 40 | max_displacement: int |
| 41 | Max displacement of Correlation |
| 42 | |
| 43 | stride1: int |
| 44 | Stride for data1 |
| 45 | |
| 46 | stride2: int |
| 47 | Stride for data2 within the neightborhood centered around data1 |
| 48 | |
| 49 | padding: int |
| 50 | Padding for correlation |
| 51 | |
| 52 | is_multiply: bool |
| 53 | operation type is either multiplication or substraction |
| 54 | |
| 55 | Returns |
| 56 | ------- |
| 57 | c_np : np.ndarray |
| 58 | 4-D with shape [batch, out_channel, out_height, out_width] |
| 59 | """ |
| 60 | # compute output's dimension |
| 61 | pad_data_height = data1.shape[2] + 2 * padding |
| 62 | pad_data_width = data1.shape[3] + 2 * padding |
| 63 | kernel_radius = (kernel_size - 1) // 2 |
| 64 | border_size = max_displacement + kernel_radius |
| 65 | out_width = (pad_data_width - border_size * 2) // stride1 |
| 66 | out_height = (pad_data_height - border_size * 2) // stride1 |
| 67 | neighborhood_grid_radius = max_displacement // stride2 |
| 68 | neighborhood_grid_width = neighborhood_grid_radius * 2 + 1 |
| 69 | out_channel = neighborhood_grid_width * neighborhood_grid_width |
| 70 | |
| 71 | out = np.zeros((data1.shape[0], out_channel, out_height, out_width)) |
| 72 | pad_data1 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width)) |
| 73 | pad_data2 = np.zeros((data1.shape[0], data1.shape[1], pad_data_height, pad_data_width)) |
| 74 | |
| 75 | pad_data1[:, :, padding : padding + data1.shape[2], padding : padding + data1.shape[3]] = data1[ |
| 76 | :, :, :, : |
| 77 | ] |
| 78 | pad_data2[:, :, padding : padding + data2.shape[2], padding : padding + data2.shape[3]] = data2[ |
| 79 | :, :, :, : |
| 80 | ] |
| 81 |