| 43 | } |
| 44 | |
| 45 | af_err af_nearest_neighbour(af_array* idx, af_array* dist, const af_array query, |
| 46 | const af_array train, const dim_t dist_dim, |
| 47 | const uint n_dist, const af_match_type dist_type) { |
| 48 | try { |
| 49 | const ArrayInfo& qInfo = getInfo(query); |
| 50 | const ArrayInfo& tInfo = getInfo(train); |
| 51 | af_dtype qType = qInfo.getType(); |
| 52 | af_dtype tType = tInfo.getType(); |
| 53 | af::dim4 qDims = qInfo.dims(); |
| 54 | af::dim4 tDims = tInfo.dims(); |
| 55 | |
| 56 | uint train_samples = (dist_dim == 0) ? 1 : 0; |
| 57 | |
| 58 | DIM_ASSERT(2, qDims[dist_dim] == tDims[dist_dim]); |
| 59 | DIM_ASSERT(2, qDims[2] == 1 && qDims[3] == 1); |
| 60 | DIM_ASSERT(3, tDims[2] == 1 && tDims[3] == 1); |
| 61 | DIM_ASSERT(4, (dist_dim == 0 || dist_dim == 1)); |
| 62 | DIM_ASSERT(5, n_dist > 0 && n_dist <= (uint)tDims[train_samples]); |
| 63 | ARG_ASSERT(5, n_dist > 0 && n_dist <= 256); |
| 64 | ARG_ASSERT(6, dist_type == AF_SAD || dist_type == AF_SSD || |
| 65 | dist_type == AF_SHD); |
| 66 | TYPE_ASSERT(qType == tType); |
| 67 | |
| 68 | // For Hamming, only u8, u16, u32 and u64 allowed. |
| 69 | af_array oIdx; |
| 70 | af_array oDist; |
| 71 | |
| 72 | if (dist_type == AF_SHD) { |
| 73 | TYPE_ASSERT(qType == u8 || qType == u16 || qType == u32 || |
| 74 | qType == u64); |
| 75 | switch (qType) { |
| 76 | case u8: |
| 77 | nearest_neighbour<uchar, uint>(&oIdx, &oDist, query, train, |
| 78 | dist_dim, n_dist, AF_SHD); |
| 79 | break; |
| 80 | case u16: |
| 81 | nearest_neighbour<ushort, uint>(&oIdx, &oDist, query, train, |
| 82 | dist_dim, n_dist, AF_SHD); |
| 83 | break; |
| 84 | case u32: |
| 85 | nearest_neighbour<uint, uint>(&oIdx, &oDist, query, train, |
| 86 | dist_dim, n_dist, AF_SHD); |
| 87 | break; |
| 88 | case u64: |
| 89 | nearest_neighbour<uintl, uint>(&oIdx, &oDist, query, train, |
| 90 | dist_dim, n_dist, AF_SHD); |
| 91 | break; |
| 92 | default: TYPE_ERROR(1, qType); |
| 93 | } |
| 94 | } else { |
| 95 | switch (qType) { |
| 96 | case f32: |
| 97 | nearest_neighbour<float, float>(&oIdx, &oDist, query, train, |
| 98 | dist_dim, n_dist, |
| 99 | dist_type); |
| 100 | break; |
| 101 | case f64: |
| 102 | nearest_neighbour<double, double>(&oIdx, &oDist, query, |
no test coverage detected