Refactored cpu implementatio and added cuda aproach.

master
shugeo 2019-10-07 17:51:07 +03:00
parent 78443ffebf
commit 6cf3a8fa9c
3 changed files with 70 additions and 21 deletions

View File

@ -32,28 +32,25 @@ namespace helpers {
// height = images->sizeAt(1), width = images->sizeAt(2)
// colors - colors for each box given
// set up color for each box as frame
auto batchSize = images->sizeAt(0);
auto height = images->sizeAt(1);
auto width = images->sizeAt(2);
auto channels = images->sizeAt(3);
auto imageList = images->allTensorsAlongDimension({1, 2, 3}); // split images by batch
auto boxList = boxes->allTensorsAlongDimension({1, 2}); // split boxes by batch
output->assign(images);
for (auto b = 0; b < imageList->size(); ++b) { // loop by batch
//auto imageList = images->allTensorsAlongDimension({1, 2, 3}); // split images by batch
// auto boxList = boxes->allTensorsAlongDimension({1, 2}); // split boxes by batch
auto colorSet = colors->allTensorsAlongDimension({1});
output->assign(images); // fill up all output with input images, then fill up boxes
for (auto b = 0; b < batchSize; ++b) { // loop by batch
// auto image = imageList->at(b);
auto box = boxList->at(b);
auto internalBoxes = box->allTensorsAlongDimension({1});
auto colorSet = colors->allTensorsAlongDimension({1});
for (auto c = 0; c < colorSet->size(); ++c) {
// box with shape
auto internalBox = internalBoxes->at(c);
auto internalBox = (*boxes)(b, {0})(c, {0});//internalBoxes->at(c);
auto color = colorSet->at(c);
auto rowStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((height - 1) * internalBox->e<float>(0)));
auto rowEnd = nd4j::math::nd4j_min(Nd4jLong (height - 1), Nd4jLong ((height - 1) * internalBox->e<float>(2)));
auto colStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((width - 1) * internalBox->e<float>(1)));
auto colEnd = nd4j::math::nd4j_min(Nd4jLong(width - 1), Nd4jLong ((width - 1) * internalBox->e<float>(3)));
auto rowStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((height - 1) * internalBox.e<float>(0)));
auto rowEnd = nd4j::math::nd4j_min(Nd4jLong (height - 1), Nd4jLong ((height - 1) * internalBox.e<float>(2)));
auto colStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((width - 1) * internalBox.e<float>(1)));
auto colEnd = nd4j::math::nd4j_min(Nd4jLong(width - 1), Nd4jLong ((width - 1) * internalBox.e<float>(3)));
for (auto y = rowStart; y <= rowEnd; y++) {
for (auto e = 0; e < color->lengthOf(); ++e) {
output->p(b, y, colStart, e, color->e(e));
@ -67,11 +64,11 @@ namespace helpers {
}
}
}
delete colorSet;
delete internalBoxes;
// delete internalBoxes;
}
delete imageList;
delete boxList;
delete colorSet;
// delete imageList;
// delete boxList;
}
}

View File

@ -24,8 +24,60 @@ namespace nd4j {
namespace ops {
namespace helpers {
template <typename T>
static __global__ void drawBoundignBoxesKernel(T const* images, Nd4jLong* imagesShape, T const* boxes,
Nd4jLong* boxesShape, T const* colors, Nd4jLong* colorsShape, T* output, Nd4jLong* outputShape) {
for (auto b = 0; b < batchSize; ++b) { // loop by batch
for (auto c = 0; c < colorSet->size(); ++c) {
// box with shape
auto internalBox = (*boxes)(b, {0})(c, {0});//internalBoxes->at(c);
auto color = colorSet->at(c);
auto rowStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((height - 1) * internalBox.e<float>(0)));
auto rowEnd = nd4j::math::nd4j_min(Nd4jLong (height - 1), Nd4jLong ((height - 1) * internalBox.e<float>(2)));
auto colStart = nd4j::math::nd4j_max(Nd4jLong (0), Nd4jLong ((width - 1) * internalBox.e<float>(1)));
auto colEnd = nd4j::math::nd4j_min(Nd4jLong(width - 1), Nd4jLong ((width - 1) * internalBox.e<float>(3)));
for (auto y = rowStart; y <= rowEnd; y++) {
for (auto e = 0; e < color->lengthOf(); ++e) {
output->p(b, y, colStart, e, color->e(e));
output->p(b, y, colEnd, e, color->e(e));
}
}
for (auto x = colStart + 1; x < colEnd; x++) {
for (auto e = 0; e < color->lengthOf(); ++e) {
output->p(b, rowStart, x, e, color->e(e));
output->p(b, rowEnd, x, e, color->e(e));
}
}
}
}
}
template <typename T>
void drawBoundingBoxesH(nd4j::LaunchContext* context, NDArray const* images, NDArray const* boxes, NDArray const* colors, NDArray* output) {
drawBoundingBoxesKernel<<<128, 256, 1024, *stream>>>(imagesBuf, imagesShape, boxesBuf, boxesShape, colorsBuf, colorsShape,
outputBuf, outputShape, );
}
void drawBoundingBoxesFunctor(nd4j::LaunchContext * context, NDArray* images, NDArray* boxes, NDArray* colors, NDArray* output) {
// images - batch of 3D images with BW (last dim = 1), RGB (last dim = 3) or RGBA (last dim = 4) channel set
// boxes - batch of 2D bounds with last dim (y_start, x_start, y_end, x_end) to compute i and j as
// floor((height - 1 ) * y_start) => rowStart, floor((height - 1) * y_end) => rowEnd
// floor((width - 1 ) * x_start) => colStart, floor((width - 1) * x_end) => colEnd
// height = images->sizeAt(1), width = images->sizeAt(2)
// colors - colors for each box given
// set up color for each box as frame
auto height = images->sizeAt(1);
auto width = images->sizeAt(2);
auto channels = images->sizeAt(3);
auto imageList = images->allTensorsAlongDimension({1, 2, 3}); // split images by batch
auto boxList = boxes->allTensorsAlongDimension({1, 2}); // split boxes by batch
auto colorSet = colors->allTensorsAlongDimension({1});
output->assign(images);
delete colorSet;
delete imageList;
delete boxList;
}
}

View File

@ -2072,8 +2072,8 @@ TEST_F(DeclarableOpsTests10, Image_DrawBoundingBoxes_1) {
ASSERT_EQ(ND4J_STATUS_OK, results->status());
auto result = results->at(0);
result->printBuffer("Bounded boxes");
expected.printBuffer("Bounded expec");
// result->printBuffer("Bounded boxes");
// expected.printBuffer("Bounded expec");
ASSERT_TRUE(expected.isSameShapeStrict(result));
ASSERT_TRUE(expected.equalsTo(result));