Commit 6294701e authored by Nianchen Deng's avatar Nianchen Deng
Browse files

sync

parent 2824f796
#include "Msl.h"
#include <time.h>
Msl::Msl() : net(nullptr) {}
bool Msl::load(const std::string &netPath) {
net = new Net();
if (net->load(netPath))
return true;
dispose();
return false;
}
void Msl::bindResources(Resource *resEncoded, Resource *resDepths, Resource *resColors) {
net->bindResource("Encoded", resEncoded);
net->bindResource("Depths", resDepths);
net->bindResource("Colors", resColors);
}
bool Msl::infer() { return net->infer(); }
void Msl::dispose() {
if (net != nullptr) {
net->dispose();
delete net;
net = nullptr;
}
}
#pragma once
#include "../utils/common.h"
#include "Net.h"
class Msl {
public:
Net *net;
Msl();
virtual bool load(const std::string &netDir);
virtual void bindResources(Resource *resEncoded, Resource *resDepths, Resource *resColors);
virtual bool infer();
virtual void dispose();
};
#pragma once
#include "../utils/common.h"
class Net {
public:
Net() : mEngine(nullptr) { }
bool load(const std::string& path);
void bindResource(const std::string& name, Resource* res);
bool dispose();
bool infer(cudaStream_t stream = nullptr, bool dumpInputOutput = false);
private:
std::shared_ptr<nv::ICudaEngine> mEngine;
std::shared_ptr<nv::IExecutionContext> mContext;
Resources mResources;
void _deserialize(const std::string& path);
std::vector<void*> _getBindings();
void _dumpInputOutput();
protected:
bool _dumpBuffer(std::ostream& os, void* deviceBuf, int index);
bool _dumpBuffer(std::ostream& os, void* deviceBuf, nv::Dims bufDims, nv::DataType dataType);
};
#include "Nmsl2.h"
#include <time.h>
Nmsl2::Nmsl2(int batchSize, int samples)
: batchSize(batchSize),
samples(samples),
resRaw1(nullptr),
resRaw2(nullptr),
fcNet1(nullptr),
fcNet2(nullptr),
catNet(nullptr) {}
bool Nmsl2::load(const std::string &netDir) {
fcNet1 = new Net();
fcNet2 = new Net();
catNet = new Net();
if (!fcNet1->load(netDir + "fc1.trt") || !fcNet2->load(netDir + "fc2.trt") ||
!catNet->load(netDir + "cat.trt"))
return false;
resRaw1 = sptr<Resource>(new CudaBuffer(batchSize * samples / 2 * sizeof(float4)));
resRaw2 = sptr<Resource>(new CudaBuffer(batchSize * samples / 2 * sizeof(float4)));
return true;
}
void Nmsl2::bindResources(Resource *resEncoded, Resource *resDepths, Resource *resColors) {
fcNet1->bindResource("Encoded", resEncoded);
fcNet1->bindResource("Raw", resRaw1.get());
fcNet2->bindResource("Encoded", resEncoded);
fcNet2->bindResource("Raw", resRaw2.get());
catNet->bindResource("Raw1", resRaw1.get());
catNet->bindResource("Raw2", resRaw2.get());
catNet->bindResource("Depths", resDepths);
catNet->bindResource("Colors", resColors);
}
bool Nmsl2::infer() {
// CudaStream stream1, stream2;
if (!fcNet1->infer())
return false;
if (!fcNet2->infer())
return false;
if (!catNet->infer())
return false;
return true;
}
void Nmsl2::dispose() {
if (fcNet1 != nullptr) {
fcNet1->dispose();
delete fcNet1;
fcNet1 = nullptr;
}
if (fcNet2 != nullptr) {
fcNet2->dispose();
delete fcNet2;
fcNet2 = nullptr;
}
if (catNet != nullptr) {
catNet->dispose();
delete catNet;
catNet = nullptr;
}
resRaw1 = nullptr;
resRaw2 = nullptr;
}
#include "Renderer.h"
#include "../utils/cuda.h"
/// Dispatch (n_rays, -)
__global__ void cu_render(glm::vec4 *o_colors, glm::vec4 *layeredColors, uint samples, uint nRays) {
glm::uvec3 idx3 = IDX3;
uint rayIdx = idx3.x;
if (rayIdx >= nRays)
return;
glm::vec4 outColor;
for (int si = samples - 1; si >= 0; --si) {
glm::vec4 c = layeredColors[rayIdx * samples + si];
outColor = outColor * (1 - c.a) + c * c.a;
}
outColor.a = 1.0f;
o_colors[idx3.x] = outColor;
}
Renderer::Renderer() {}
void Renderer::render(sptr<CudaArray<glm::vec4>> o_colors,
sptr<CudaArray<glm::vec4>> layeredColors) {
dim3 blkSize(1024);
dim3 grdSize(ceilDiv(o_colors->n(), blkSize.x));
CU_INVOKE(cu_render)
(*o_colors, *layeredColors, layeredColors->n() / o_colors->n(), o_colors->n());
CHECK_EX(cudaGetLastError());
}
\ No newline at end of file
#pragma once
#include "../utils/common.h"
class Renderer {
public:
Renderer();
/**
* @brief
*
* @param o_colors
* @param layeredColors
*/
void render(sptr<CudaArray<glm::vec4>> o_colors, sptr<CudaArray<glm::vec4>> layeredColors);
};
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment