geom: fix interaction of qtbt and octree geometry quantisation

The geometry quantisation assumed that octree nodes were cubes.  The
addition of qt/bt partitioning modes breaks this premise.  This commit
addresses the following:

 - Only the to-be-coded bits, which may differ for each component, are
   quantised, rather than quantising all bits according to the tree
   depth (which would result in quantising already coded bits).

 - The occupancy mask (determining the non-coded occupancy bits) must
   be recomputed for each node.  Since the effect of quantisation is
   to prune back each sub-tree from the leaves, some nodes that would
   have had qt/bt partitions will not be coded, leaving an unequal
   reconstruction.  The solution is to terminate the shorter leaves
   earlier and to recompute the partition.

tmc3/geometry_octree_decoder.cpp

@@ -475,7 +475,7 @@ GeometryOctreeDecoder::decodeDirectPosition(
 // Helper to inverse quantise positions
 
 Vec3<uint32_t>
-invQuantPosition(int qp, int quantBitMask, const Vec3<uint32_t>& pos)
+invQuantPosition(int qp, Vec3<uint32_t> quantMasks, const Vec3<uint32_t>& pos)
 {
   // pos represents the position within the coded tree as follows:
   //   |pppppqqqqqq|00
@@ -489,7 +489,7 @@ invQuantPosition(int qp, int quantBitMask, const Vec3<uint32_t>& pos)
   int shiftBits = (qp - 4) / 6;
   Vec3<uint32_t> recon;
   for (int k = 0; k < 3; k++) {
-    int posQuant = pos[k] & (quantBitMask >> shiftBits);
+    int posQuant = pos[k] & (quantMasks[k] >> shiftBits);
     recon[k] = (pos[k] ^ posQuant) << shiftBits;
     recon[k] |= quantizer.scale(posQuant);
   }
@@ -570,15 +570,12 @@ decodeGeometryOctree(
 
   int sliceQp = gps.geom_base_qp + gbh.geom_slice_qp_offset;
   int numLvlsUntilQpOffset = -1;
-  int posQuantBits = 0;
+  Vec3<uint32_t> posQuantBitMasks = 0xffffffff;
 
   if (gbh.geom_octree_qp_offset_enabled_flag)
     numLvlsUntilQpOffset = gbh.geom_octree_qp_offset_depth;
-  else if (gps.geom_scaling_enabled_flag) {
+  else if (gps.geom_scaling_enabled_flag)
     node00.qp = sliceQp;
-    // determine the mask of LSBs used in quantisation
-    posQuantBits = (1 << (nodeMaxDimLog2 - numLvlsUntilQpOffset)) - 1;
-  }
 
   for (; !fifo.empty(); fifo.pop_front()) {
     if (fifo.begin() == fifoCurrLvlEnd) {
@@ -596,7 +593,7 @@ decodeGeometryOctree(
       if (maxNumImplicitQtbtBeforeOt)
         maxNumImplicitQtbtBeforeOt--;
 
-      // if all dimensions have same size, then performing octree for remaining nodes
+      // if all dimensions have same size, then use octree for remaining nodes
       if (
         nodeSizeLog2[0] == nodeSizeLog2[1]
         && nodeSizeLog2[1] == nodeSizeLog2[2])
@@ -606,8 +603,6 @@ decodeGeometryOctree(
       childSizeLog2 = implicitQtBtDecision(
         nodeSizeLog2, maxNumImplicitQtbtBeforeOt, minSizeImplicitQtbt);
 
-      occupancySkip = nonSplitQtBtAxes(nodeSizeLog2, childSizeLog2);
-
       nodeMaxDimLog2--;
       numNodesNextLvl = 0;
       occupancyAtlasOrigin = 0xffffffff;
@@ -622,7 +617,12 @@ decodeGeometryOctree(
       if (nodeMaxDimLog2 == minNodeSizeLog2)
         break;
 
+      // record the node size when quantisation is signalled -- all subsequnt
+      // coded occupancy bits are quantised
       numLvlsUntilQpOffset--;
+      if (!numLvlsUntilQpOffset)
+        for (int k = 0; k < 3; k++)
+          posQuantBitMasks[k] = (1 << nodeSizeLog2[k]) - 1;
     }
 
     PCCOctree3Node& node0 = fifo.front();
@@ -635,6 +635,16 @@ decodeGeometryOctree(
     auto effectiveNodeSizeLog2 = nodeSizeLog2 - shiftBits;
     auto effectiveChildSizeLog2 = childSizeLog2 - shiftBits;
 
+    // todo(??): the following needs to be reviewed, it is added to make
+    // quantisation work with qtbt.
+    Vec3<int> actualNodeSizeLog2, actualChildSizeLog2;
+    for (int k = 0; k < 3; k++) {
+      actualNodeSizeLog2[k] = std::max(nodeSizeLog2[k], shiftBits);
+      actualChildSizeLog2[k] = std::max(childSizeLog2[k], shiftBits);
+    }
+    // todo(??): atlasShift may be wrong too
+    occupancySkip = nonSplitQtBtAxes(actualNodeSizeLog2, actualChildSizeLog2);
+
     int occupancyAdjacencyGt0 = 0;
     int occupancyAdjacencyGt1 = 0;
     int occupancyAdjacencyUnocc = 0;
@@ -710,7 +720,7 @@ decodeGeometryOctree(
                            (node0.pos[1] << !(occupancySkip & 2)) + y,
                            (node0.pos[2] << !(occupancySkip & 1)) + z};
 
-        pos = invQuantPosition(node0.qp, posQuantBits, pos);
+        pos = invQuantPosition(node0.qp, posQuantBitMasks, pos);
         const Vec3<double> point(pos[0], pos[1], pos[2]);
 
         for (int i = 0; i < numPoints; ++i)
@@ -747,7 +757,7 @@ decodeGeometryOctree(
             point[k] += child.pos[k] << shift;
           }
 
-          point = invQuantPosition(node0.qp, posQuantBits, point);
+          point = invQuantPosition(node0.qp, posQuantBitMasks, point);
           pointCloud[processedPointCount++] =
             Vec3<double>(point[0], point[1], point[2]);
         }
@@ -784,7 +794,7 @@ decodeGeometryOctree(
       int quantRemovedBits = (node.qp - 4) / 6;
       for (int k = 0; k < 3; k++)
         node.pos[k] <<= nodeSizeLog2[k] - quantRemovedBits;
-      node.pos = invQuantPosition(node.qp, posQuantBits, node.pos);
+      node.pos = invQuantPosition(node.qp, posQuantBitMasks, node.pos);
     }
     *nodesRemaining = std::move(fifo);
     return;

tmc3/geometry_octree_encoder.cpp

@@ -487,17 +487,19 @@ calculateNodeQps(int baseQp, It nodesBegin, It nodesEnd)
 
 void
 geometryQuantization(
-  PCCPointSet3& pointCloud, PCCOctree3Node& node, int nodeSizeLog2)
+  PCCPointSet3& pointCloud, PCCOctree3Node& node, Vec3<int> nodeSizeLog2)
 {
   QuantizerGeom quantizer = QuantizerGeom(node.qp);
   int qpShift = (node.qp - 4) / 6;
-  int quantBitsMask = (1 << nodeSizeLog2) - 1;
-  Vec3<uint32_t> end_pos_quant = ((1 << nodeSizeLog2) >> qpShift) - 1;
-  for (int i = node.start; i < node.end; i++) {
-    for (int k = 0; k < 3; k++) {
+
+  for (int k = 0; k < 3; k++) {
+    int quantBitsMask = (1 << nodeSizeLog2[k]) - 1;
+    uint32_t clipMax = ((1 << nodeSizeLog2[k]) >> qpShift) - 1;
+
+    for (int i = node.start; i < node.end; i++) {
       uint32_t pos = uint32_t(pointCloud[i][k]);
       uint32_t quantPos = quantizer.quantize(pos & quantBitsMask);
-      quantPos = PCCClip(quantPos, 0, end_pos_quant[k]);
+      quantPos = PCCClip(quantPos, 0, clipMax);
 
       // NB: this representation is: |pppppp00qqq|, which isn't the
       // same used by the decoder:  (|ppppppqqq|00)
@@ -510,13 +512,14 @@ geometryQuantization(
 
 void
 geometryScale(
-  PCCPointSet3& pointCloud, PCCOctree3Node& node, int quantNodeMaxDimLog2)
+  PCCPointSet3& pointCloud, PCCOctree3Node& node, Vec3<int> quantNodeSizeLog2)
 {
   QuantizerGeom quantizer = QuantizerGeom(node.qp);
   int qpShift = (node.qp - 4) / 6;
-  int quantBitsMask = (1 << quantNodeMaxDimLog2) - 1;
-  for (int i = node.start; i < node.end; i++) {
-    for (int k = 0; k < 3; k++) {
+
+  for (int k = 0; k < 3; k++) {
+    int quantBitsMask = (1 << quantNodeSizeLog2[k]) - 1;
+    for (int i = node.start; i < node.end; i++) {
       uint32_t pos = uint32_t(pointCloud[i][k]);
       uint32_t quantPos = pos & quantBitsMask;
       pointCloud[i][k] = (pos & ~quantBitsMask) | quantizer.scale(quantPos);
@@ -649,7 +652,7 @@ encodeGeometryOctree(
   Vec3<uint32_t> occupancyAtlasOrigin(0xffffffff);
 
   // the node size where quantisation is performed
-  int quantNodeMaxDimLog2 = 0;
+  Vec3<int> quantNodeSizeLog2 = 0;
   int numLvlsUntilQuantization = -1;
   if (gps.geom_scaling_enabled_flag) {
     numLvlsUntilQuantization = 0;
@@ -661,7 +664,7 @@ encodeGeometryOctree(
 
   // applied to the root node, just set the node qp
   if (numLvlsUntilQuantization == 0) {
-    quantNodeMaxDimLog2 = nodeMaxDimLog2;
+    quantNodeSizeLog2 = nodeSizeLog2;
     node00.qp = sliceQp;
   }
 
@@ -694,7 +697,6 @@ encodeGeometryOctree(
         nodeSizeLog2, maxNumImplicitQtbtBeforeOt, minSizeImplicitQtbt);
 
       pointSortMask = qtBtChildSize(nodeSizeLog2, childSizeLog2);
-      occupancySkip = nonSplitQtBtAxes(nodeSizeLog2, childSizeLog2);
 
       nodeMaxDimLog2--;
       encoder.beginOctreeLevel();
@@ -706,7 +708,7 @@ encodeGeometryOctree(
       // determing a per node QP at the appropriate level
       // NB: this has no effect here if geom_octree_qp_offset_depth=0
       if (--numLvlsUntilQuantization == 0) {
-        quantNodeMaxDimLog2 = nodeMaxDimLog2;
+        quantNodeSizeLog2 = nodeSizeLog2;
         calculateNodeQps(gps.geom_base_qp, fifo.begin(), fifoCurrLvlEnd);
       }
     }
@@ -723,8 +725,18 @@ encodeGeometryOctree(
     auto effectiveNodeSizeLog2 = nodeSizeLog2 - shiftBits;
     auto effectiveChildSizeLog2 = childSizeLog2 - shiftBits;
 
+    // todo(??): the following needs to be reviewed, it is added to make
+    // quantisation work with qtbt.
+    Vec3<int> actualNodeSizeLog2, actualChildSizeLog2;
+    for (int k = 0; k < 3; k++) {
+      actualNodeSizeLog2[k] = std::max(nodeSizeLog2[k], shiftBits);
+      actualChildSizeLog2[k] = std::max(childSizeLog2[k], shiftBits);
+    }
+    // todo(??): atlasShift may be wrong too
+    occupancySkip = nonSplitQtBtAxes(actualNodeSizeLog2, actualChildSizeLog2);
+
     if (numLvlsUntilQuantization == 0) {
-      geometryQuantization(pointCloud, node0, quantNodeMaxDimLog2);
+      geometryQuantization(pointCloud, node0, quantNodeSizeLog2);
       if (gps.geom_unique_points_flag)
         checkDuplicatePoints(pointCloud, node0, pointIdxToDmIdx);
     }
@@ -805,7 +817,7 @@ encodeGeometryOctree(
       int childStart = node0.start;
 
       // inverse quantise any quantised positions
-      geometryScale(pointCloud, node0, quantNodeMaxDimLog2);
+      geometryScale(pointCloud, node0, quantNodeSizeLog2);
 
       for (int i = 0; i < 8; i++) {
         if (!childCounts[i]) {
@@ -875,7 +887,7 @@ encodeGeometryOctree(
 
         if (directModeUsed) {
           // inverse quantise any quantised positions
-          geometryScale(pointCloud, node0, quantNodeMaxDimLog2);
+          geometryScale(pointCloud, node0, quantNodeSizeLog2);
 
           // point reordering to match decoder's order
           for (auto idx = child.start; idx < child.end; idx++)
@@ -912,7 +924,7 @@ encodeGeometryOctree(
     for (auto& node : fifo) {
       for (int k = 0; k < 3; k++)
         node.pos[k] <<= nodeSizeLog2[k];
-      geometryScale(pointCloud, node, quantNodeMaxDimLog2);
+      geometryScale(pointCloud, node, quantNodeSizeLog2);
     }
     *nodesRemaining = std::move(fifo);
     return;