wehub-resource-sync
80 lines
2.8 KiB
Plaintext
80 lines
2.8 KiB
Plaintext
layout(set = 0, binding = 0) writeonly buffer outputBuffer{
|
|
FLOAT4 data[];
|
|
}uOutput;
|
|
|
|
layout(set = 0, binding = 1) readonly buffer inputBuffer{
|
|
FLOAT4 data[];
|
|
}uInput;
|
|
|
|
layout(set = 0, binding = 2) readonly buffer roiBuffer{
|
|
FLOAT data[];
|
|
}uRoI;
|
|
|
|
layout(set = 0, binding = 3) uniform constBuffer
|
|
{
|
|
ivec4 inputImgSize;
|
|
ivec4 outputImgSize;
|
|
float spatialScale;
|
|
}uRoiPoolingParam;
|
|
|
|
layout(local_size_x = 8, local_size_y = 8) in;
|
|
|
|
void main()
|
|
{
|
|
ivec3 pos = ivec3(gl_GlobalInvocationID);
|
|
ivec3 outputImgSize = uRoiPoolingParam.outputImgSize.xyz;
|
|
|
|
if(pos.x < outputImgSize.x && pos.y < outputImgSize.y)
|
|
{
|
|
ivec3 inputImgSize = uRoiPoolingParam.inputImgSize.xyz;
|
|
int roiBatchIndex = pos.z % uRoiPoolingParam.outputImgSize.w;
|
|
int inputZIndex = pos.z / uRoiPoolingParam.outputImgSize.w;
|
|
// 0, xmin, ymin, xmax, ymax
|
|
FLOAT roiData0 = uRoI.data[roiBatchIndex * 5 + 0];
|
|
FLOAT roiData1 = uRoI.data[roiBatchIndex * 5 + 1];
|
|
FLOAT roiData2 = uRoI.data[roiBatchIndex * 5 + 2];
|
|
FLOAT roiData3 = uRoI.data[roiBatchIndex * 5 + 3];
|
|
FLOAT roiData4 = uRoI.data[roiBatchIndex * 5 + 4];
|
|
int inputBatch = int(roiData0);
|
|
int x1 = int(round(FLOAT(roiData1) * FLOAT(uRoiPoolingParam.spatialScale)));
|
|
int y1 = int(round(FLOAT(roiData2) * FLOAT(uRoiPoolingParam.spatialScale)));
|
|
int x2 = int(round(FLOAT(roiData3) * FLOAT(uRoiPoolingParam.spatialScale)));
|
|
int y2 = int(round(FLOAT(roiData4) * FLOAT(uRoiPoolingParam.spatialScale)));
|
|
|
|
int roiW = max(x2 - x1 + 1, 1);
|
|
int roiH = max(y2 - y1 + 1, 1);
|
|
FLOAT binSizeW = FLOAT(roiW) / FLOAT(outputImgSize.x);
|
|
FLOAT binSizeH = FLOAT(roiH) / FLOAT(outputImgSize.y);
|
|
|
|
int wStart = clamp(x1 + int(floor(FLOAT(pos.x) * binSizeW)), 0, inputImgSize.x);
|
|
int wEnd = clamp(x1 + int(ceil(FLOAT(pos.x + 1) * binSizeW)), 0, inputImgSize.x);
|
|
int hStart = clamp(y1 + int(floor(FLOAT(pos.y) * binSizeH)), 0, inputImgSize.y);
|
|
int hEnd = clamp(y1 + int(ceil(FLOAT(pos.y + 1) * binSizeH)), 0, inputImgSize.y);
|
|
|
|
bool isEmpty = (hEnd <= hStart) || (wEnd <= wStart);
|
|
|
|
int inputBasicOffset = inputBatch * uRoiPoolingParam.inputImgSize.x * uRoiPoolingParam.inputImgSize.y
|
|
+ inputZIndex * uRoiPoolingParam.inputImgSize.x * uRoiPoolingParam.inputImgSize.y * uRoiPoolingParam.inputImgSize.w
|
|
;
|
|
int inputStartOffset = wStart
|
|
+ hStart * uRoiPoolingParam.inputImgSize.x
|
|
+ inputBasicOffset;
|
|
FLOAT4 res = isEmpty ? FLOAT4(0.0) : uInput.data[inputStartOffset];
|
|
for(int i = hStart; i < hEnd; ++i)
|
|
{
|
|
for(int j = wStart; j < wEnd; ++j)
|
|
{
|
|
int inputOffset = j
|
|
+ i * uRoiPoolingParam.inputImgSize.x
|
|
+ inputBasicOffset;
|
|
res = max(res, uInput.data[inputOffset]);
|
|
}
|
|
}
|
|
uOutput.data[0
|
|
+ pos.x
|
|
+ pos.y * uRoiPoolingParam.outputImgSize.x
|
|
+ pos.z * uRoiPoolingParam.outputImgSize.x * uRoiPoolingParam.outputImgSize.y
|
|
] = res;
|
|
}
|
|
}
|