| 99 | // (since this function is called recursively, using big input sizes may lead to stack overflow) |
| 100 | template<typename ST, typename DT> |
| 101 | inline void LabelComponent(RawBufferType &tmpVec, RawBufferType &dstVec, const RawBufferType &srcVec, |
| 102 | const long4_16a &tmpStrides, const long4_16a &dstStrides, const long4_16a &srcStrides, |
| 103 | const long4_16a &shape, const long4_16a &curCoord, ST value, DT label) |
| 104 | { |
| 105 | if (util::ValueAt<U8>(tmpVec, tmpStrides, curCoord) == 1) |
| 106 | { |
| 107 | return; // The element was already labeled, skip it |
| 108 | } |
| 109 | if (value != util::ValueAt<ST>(srcVec, srcStrides, curCoord)) |
| 110 | { |
| 111 | return; // The element is not in the same labeled region, skip it |
| 112 | } |
| 113 | |
| 114 | // Set element label in dstVec and mark it as labeled in tmpVec |
| 115 | util::ValueAt<DT>(dstVec, dstStrides, curCoord) = label; |
| 116 | util::ValueAt<U8>(tmpVec, tmpStrides, curCoord) = 1; |
| 117 | |
| 118 | // For each neighbor, recursively call label component to label each neighbor |
| 119 | if (curCoord.y > 0) |
| 120 | { |
| 121 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 122 | long4_16a{curCoord.x, curCoord.y - 1, curCoord.z, curCoord.w}, value, label); |
| 123 | } |
| 124 | if (curCoord.y < shape.y - 1) |
| 125 | { |
| 126 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 127 | long4_16a{curCoord.x, curCoord.y + 1, curCoord.z, curCoord.w}, value, label); |
| 128 | } |
| 129 | if (curCoord.z > 0) |
| 130 | { |
| 131 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 132 | long4_16a{curCoord.x, curCoord.y, curCoord.z - 1, curCoord.w}, value, label); |
| 133 | } |
| 134 | if (curCoord.z < shape.z - 1) |
| 135 | { |
| 136 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 137 | long4_16a{curCoord.x, curCoord.y, curCoord.z + 1, curCoord.w}, value, label); |
| 138 | } |
| 139 | if (curCoord.w > 0) |
| 140 | { |
| 141 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 142 | long4_16a{curCoord.x, curCoord.y, curCoord.z, curCoord.w - 1}, value, label); |
| 143 | } |
| 144 | if (curCoord.w < shape.w - 1) |
| 145 | { |
| 146 | LabelComponent(tmpVec, dstVec, srcVec, tmpStrides, dstStrides, srcStrides, shape, |
| 147 | long4_16a{curCoord.x, curCoord.y, curCoord.z, curCoord.w + 1}, value, label); |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | // Label N volumes in NDHW tensor stored in srcVec yielding dstVec, with corresponding srcStrides/dstStrides |
| 152 | // - ST is the source type, the data type of the input tensor in srcVec |