szymonmaszke / torchlambda

Lightweight tool to deploy PyTorch models to AWS Lambda
MIT License
125 stars 3 forks source link

Semantic segmentation prediction in output.json #3

Open sagarr opened 4 years ago

sagarr commented 4 years ago

Hi,

Is it possible to return base64 encoded image segmentation result in output.json?

This is the code I came up with after googling

auto output = module.forward({normalized_tensor}).toTensor();
torch::Tensor outputTensor = torch::sigmoid(output.detach().to(torch::kCPU));
outputTensor= torch::where(outputTensor >= 0.75, torch::ones(1), torch::zeros(1));
outputTensor = outputTensor.permute({ 0, 2, 3, 1 });
outputTensor = outputTensor.squeeze(0).detach();
outputTensor = outputTensor.mul(255).clamp(0, 255).to(torch::kU8);
outputTensor = outputTensor.to(torch::kCPU);
const auto outBuffer = outputTensor.data_ptr();
// TODO encode to base64 and return as json string

Not sure how can I return base64 encoded string back?

PS: Im not familiar with C++

szymonmaszke commented 4 years ago

Yes it should be fairly easy with .yaml settings. Will add a comment with it tomorrow.

BTW. Usually you don't need to delve into C++ with this tool.

szymonmaszke commented 4 years ago

Update: base64 return type is not currently supported, you would have to create C++ on your own based on torchlambda template generated setup.

szymonmaszke commented 4 years ago

For now, you can see how to modify C++ source code in this issue, transforming to base64 needs some knowledge of C++ and AWS API & PyTorch C++ API though.