Skip to content

Commit 9f2ea57

Browse files
serizbaserizba
committed
Added support to unknown shapes
1 parent 98312ce commit 9f2ea57

File tree

7 files changed

+122
-7
lines changed

7 files changed

+122
-7
lines changed

‎examples/coco/CMakeLists.txt

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
cmake_minimum_required(VERSION 3.10)
2+
project(example)
3+
4+
set(CMAKE_CXX_STANDARD 17)
5+
add_executable(example main.cpp ../../src/Model.cpp ../../src/Tensor.cpp)
6+
find_package( OpenCV REQUIRED )
7+
target_include_directories(example PRIVATE ../../include)
8+
target_link_libraries (example -ltensorflow ${OpenCV_LIBS})

‎examples/coco/main.cpp

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
//
2+
// Created by sergio on 16/05/19.
3+
//
4+
5+
#include "../../include/Model.h"
6+
#include "../../include/Tensor.h"
7+
#include <opencv2/opencv.hpp>
8+
#include <numeric>
9+
#include <iomanip>
10+
11+
int main() {
12+
Model model("../ssd_inception/frozen_inference_graph.pb");
13+
// Model model("../model.pb");
14+
// model.init();
15+
auto outNames1 = new Tensor(model, "num_detections");
16+
auto outNames2 = new Tensor(model, "detection_scores");
17+
auto outNames3 = new Tensor(model, "detection_boxes");
18+
auto outNames4 = new Tensor(model, "detection_classes");
19+
20+
auto inpName = new Tensor(model, "image_tensor");
21+
22+
23+
// Read image
24+
cv::Mat img, inp;
25+
img = cv::imread("../test.jpg", CV_LOAD_IMAGE_COLOR);
26+
27+
int rows = img.rows;
28+
int cols = img.cols;
29+
30+
cv::resize(img, inp, cv::Size(300, 300));
31+
cv::cvtColor(inp, inp, CV_BGR2RGB);
32+
33+
// Put image in Tensor
34+
std::vector<uint8_t > img_data;
35+
img_data.assign(inp.data, inp.data + inp.total() * inp.channels());
36+
inpName->set_data(img_data, {1, 300, 300, 3});
37+
38+
model.run(inpName, {outNames1, outNames2, outNames3, outNames4});
39+
40+
// Visualize detected bounding boxes.
41+
int num_detections = (int)outNames1->get_data<float>()[0];
42+
for (int i=0; i<num_detections; i++) {
43+
int classId = (int)outNames4->get_data<float>()[i];
44+
float score = outNames2->get_data<float>()[i];
45+
auto bbox_data = outNames3->get_data<float>();
46+
std::vector<float> bbox = {bbox_data[i*4], bbox_data[i*4+1], bbox_data[i*4+2], bbox_data[i*4+3]};
47+
if (score > 0.3) {
48+
float x = bbox[1] * cols;
49+
float y = bbox[0] * rows;
50+
float right = bbox[3] * cols;
51+
float bottom = bbox[2] * rows;
52+
53+
cv::rectangle(img, {(int)x, (int)y}, {(int)right, (int)bottom}, {125, 255, 51}, 2);
54+
}
55+
}
56+
57+
cv::imshow("Result", img);
58+
cv::imwrite("test-result.jpg", img);
59+
cv::waitKey(0);
60+
}

‎examples/coco/test-result.jpg

153 KB
Loading

‎examples/coco/test.jpg

115 KB
Loading

‎include/Tensor.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class Tensor {
2323
void clean();
2424

2525
template<typename T> void set_data(std::vector<T> new_data);
26+
template<typename T> void set_data(std::vector<T> new_data, const std::vector<int64_t>& new_shape);
2627
template<typename T> std::vector<T> get_data();
2728

2829

@@ -38,6 +39,7 @@ class Tensor {
3839
// Aux functions
3940
void error_check(bool condition, const std::string& error);
4041
template <typename T> static TF_DataType deduce_type();
42+
void deduce_shape(const Model& model);
4143

4244
public:
4345
friend class Model;

‎src/Model.cpp

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ Model::Model(const std::string& model_filename) {
1313

1414
// Create the session.
1515
TF_SessionOptions* sess_opts = TF_NewSessionOptions();
16+
1617
this->session = TF_NewSession(this->graph, sess_opts, this->status);
1718
TF_DeleteSessionOptions(sess_opts);
1819

@@ -55,12 +56,15 @@ void Model::init() {
5556

5657
void Model::save(const std::string &ckpt) {
5758
// Encode file_name to tensor
58-
size_t size = 8 + TF_StringEncodedSize(ckpt.size());
59+
size_t size = 8 + TF_StringEncodedSize(ckpt.length());
5960
TF_Tensor* t = TF_AllocateTensor(TF_STRING, nullptr, 0, size);
6061
char* data = static_cast<char *>(TF_TensorData(t));
6162
for (int i=0; i<8; i++) {data[i]=0;}
6263
TF_StringEncode(ckpt.c_str(), ckpt.size(), data + 8, size - 8, status);
6364

65+
memset(data, 0, 8); // 8-byte offset of first string.
66+
TF_StringEncode(ckpt.c_str(), ckpt.length(), (char*)(data + 8), size - 8, status);
67+
6468
// Check errors
6569
if (!this->status_check(false)) {
6670
TF_DeleteTensor(t);
@@ -81,7 +85,6 @@ void Model::save(const std::string &ckpt) {
8185
}
8286

8387

84-
8588
TF_SessionRun(this->session, nullptr, inputs, input_values, 1, nullptr, nullptr, 0, restore_op, 1, nullptr, this->status);
8689
TF_DeleteTensor(t);
8790

@@ -208,6 +211,7 @@ void Model::run(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>&
208211
for (int i=0; i<outputs.size(); i++) {
209212
outputs[i]->val = ov[i];
210213
outputs[i]->flag = 1;
214+
outputs[i]->deduce_shape(*this);
211215
}
212216

213217
// Mark input as empty

‎src/Tensor.cpp

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
#include <utility>
2+
13
//
24
// Created by sergio on 13/05/19.
35
//
@@ -19,9 +21,6 @@ Tensor::Tensor(const Model& model, const std::string& operation) {
1921
// Get number of dimensions
2022
int n_dims = TF_GraphGetTensorNumDims(model.graph, this->op, model.status);
2123

22-
// Dimension is not known
23-
error_check(n_dims != -1, "Shape of tensors must be known");
24-
2524
// DataType
2625
this->type = TF_OperationOutputType(this->op);
2726

@@ -37,8 +36,7 @@ Tensor::Tensor(const Model& model, const std::string& operation) {
3736
this->shape = std::vector<int64_t>(dims, dims + n_dims);
3837

3938
// Only one dimension can be unknown using this constructor
40-
error_check(std::count(this->shape.begin(), this->shape.end(), -1) <= 1,
41-
"At most one dimension can be unknown");
39+
// error_check(std::count(this->shape.begin(), this->shape.end(), -1) <= 1, "At most one dimension can be unknown");
4240
}
4341

4442
this->flag = 0;
@@ -82,6 +80,12 @@ void Tensor::set_data(std::vector<T> new_data) {
8280
// Check type
8381
this->error_check(deduce_type<T>() == this->type, "Provided type is different from Tensor expected type");
8482

83+
// Dimensions must be known
84+
this->error_check(!this->shape.empty(), "Shape of the input Tensor is not known, please provide a shape");
85+
86+
// At most one dimension can be unknown
87+
this->error_check(std::count(this->shape.begin(), this->shape.end(), -1) >= -1, "At most one dimension can be unknown, please provide a shape");
88+
8589
// Check number of elements
8690
auto exp_size = std::abs(std::accumulate(this->shape.begin(), this->shape.end(), 1, std::multiplies<>()));
8791

@@ -107,6 +111,17 @@ void Tensor::set_data(std::vector<T> new_data) {
107111
this->flag = 1;
108112
}
109113

114+
template<typename T> void Tensor::set_data(std::vector<T> new_data, const std::vector<int64_t>& new_shape) {
115+
116+
this->error_check(this->shape.empty() || this->shape.size() == new_shape.size(), "Provided shape has different number of dimensions");
117+
auto old_shape = this->shape;
118+
119+
this->shape = new_shape;
120+
this->set_data(new_data);
121+
122+
this->shape = old_shape;
123+
}
124+
110125
template<typename T>
111126
std::vector<T> Tensor::get_data() {
112127

@@ -159,6 +174,20 @@ TF_DataType Tensor::deduce_type() {
159174
return TF_UINT64;
160175
}
161176

177+
void Tensor::deduce_shape(const Model& model) {
178+
// Get number of dimensions
179+
int n_dims = TF_NumDims(this->val);
180+
181+
// If is not a scalar
182+
if (n_dims > 0) {
183+
// Get dimensions
184+
this->shape = std::vector<int64_t>(n_dims, -1);
185+
for (int i=0; i<n_dims; i++) {
186+
this->shape[i] = TF_Dim(this->val, i);
187+
}
188+
}
189+
}
190+
162191

163192
// VALID deduce_type TEMPLATES
164193
template TF_DataType Tensor::deduce_type<float>();
@@ -199,3 +228,15 @@ template void Tensor::set_data<uint16_t>(std::vector<uint16_t> new_data);
199228
template void Tensor::set_data<uint32_t>(std::vector<uint32_t> new_data);
200229
template void Tensor::set_data<uint64_t>(std::vector<uint64_t> new_data);
201230

231+
// VALID set_data TEMPLATES
232+
template void Tensor::set_data<float>(std::vector<float> new_data, const std::vector<int64_t>& new_shape);
233+
template void Tensor::set_data<double>(std::vector<double> new_data, const std::vector<int64_t>& new_shape);
234+
//template void Tensor::set_data<bool>(std::vector<bool> new_data, const std::vector<int64_t>& new_shape);
235+
template void Tensor::set_data<int8_t>(std::vector<int8_t> new_data, const std::vector<int64_t>& new_shape);
236+
template void Tensor::set_data<int16_t>(std::vector<int16_t> new_data, const std::vector<int64_t>& new_shape);
237+
template void Tensor::set_data<int32_t>(std::vector<int32_t> new_data, const std::vector<int64_t>& new_shape);
238+
template void Tensor::set_data<int64_t>(std::vector<int64_t> new_data, const std::vector<int64_t>& new_shape);
239+
template void Tensor::set_data<uint8_t>(std::vector<uint8_t> new_data, const std::vector<int64_t>& new_shape);
240+
template void Tensor::set_data<uint16_t>(std::vector<uint16_t> new_data, const std::vector<int64_t>& new_shape);
241+
template void Tensor::set_data<uint32_t>(std::vector<uint32_t> new_data, const std::vector<int64_t>& new_shape);
242+
template void Tensor::set_data<uint64_t>(std::vector<uint64_t> new_data, const std::vector<int64_t>& new_shape);

0 commit comments

Comments
 (0)