AttributeEncoder.cpp 36.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
/* The copyright in this software is being made available under the BSD
 * Licence, included below.  This software may be subject to other third
 * party and contributor rights, including patent rights, and no such
 * rights are granted under this licence.
 *
 * Copyright (c) 2017-2018, ISO/IEC
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 *
 * * Neither the name of the ISO/IEC nor the names of its contributors
 *   may be used to endorse or promote products derived from this
 *   software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "AttributeEncoder.h"
37

David Flynn's avatar
David Flynn committed
38
#include "ArithmeticCodec.h"
39
#include "DualLutCoder.h"
40
#include "constants.h"
41
#include "entropy.h"
42
#include "quantization.h"
43
#include "RAHT.h"
44
#include "FixedPoint.h"
45

46
// todo(df): promote to per-attribute encoder parameter
47
48
static const double kAttrPredLambdaR = 0.01;
static const double kAttrPredLambdaC = 0.01;
49

50
51
52
53
54
namespace pcc {
//============================================================================
// An encapsulation of the entropy coding methods used in attribute coding

struct PCCResidualsEncoder {
55
56
57
58
  EntropyEncoder arithmeticEncoder;
  StaticBitModel binaryModel0;
  AdaptiveBitModel binaryModelDiff[7];
  AdaptiveBitModel binaryModelIsZero[7];
59
  AdaptiveBitModel ctxPredMode[2];
60
61
  AdaptiveBitModel ctxZeroCnt[3];
  AdaptiveBitModel binaryModelIsOne[7];
62
  DualLutCoder<false> symbolCoder[2];
63

64
  void start(const SequenceParameterSet& sps, int numPoints);
65
  int stop();
66
  void encodePredMode(int value, int max);
67
  void encodeZeroCnt(int value, int max);
68
  void encodeSymbol(uint32_t value, int k1, int k2, int k3);
69
70
  void encode(uint32_t value0, uint32_t value1, uint32_t value2);
  void encode(uint32_t value);
71
72
73
74
};

//----------------------------------------------------------------------------

75
void
76
PCCResidualsEncoder::start(const SequenceParameterSet& sps, int pointCount)
77
{
78
79
  // todo(df): remove estimate when arithmetic codec is replaced
  int maxAcBufLen = pointCount * 3 * 2 + 1024;
80
  arithmeticEncoder.setBuffer(maxAcBufLen, nullptr);
81
  arithmeticEncoder.enableBypassStream(sps.cabac_bypass_stream_enabled_flag);
82
  arithmeticEncoder.start();
83
84
85
86
}

//----------------------------------------------------------------------------

87
int
88
89
PCCResidualsEncoder::stop()
{
90
  return arithmeticEncoder.stop();
91
92
93
94
}

//----------------------------------------------------------------------------

95
void
96
PCCResidualsEncoder::encodePredMode(int mode, int maxMode)
97
{
98
99
100
101
102
103
104
105
106
107
108
109
110
  // max = 0 => no direct predictors are used
  if (maxMode == 0)
    return;

  int ctxIdx = 0;
  for (int i = 0; i < mode; i++) {
    arithmeticEncoder.encode(1, ctxPredMode[ctxIdx]);
    ctxIdx = 1;
  }

  // Truncated unary
  if (mode != maxMode)
    arithmeticEncoder.encode(0, ctxPredMode[ctxIdx]);
111
112
113
114
}

//----------------------------------------------------------------------------

115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
void
PCCResidualsEncoder::encodeZeroCnt(int mode, int maxMode)
{
  // max = 0 => no direct predictors are used
  if (maxMode == 0)
    return;

  int ctxIdx = 0;
  for (int i = 0; i < mode; i++) {
    arithmeticEncoder.encode(1, ctxZeroCnt[ctxIdx]);
    ctxIdx = (ctxIdx == 0 ? 1 : 2);
  }

  // Truncated unary
  if (mode != maxMode)
    arithmeticEncoder.encode(0, ctxZeroCnt[ctxIdx]);
}

//----------------------------------------------------------------------------

135
void
136
PCCResidualsEncoder::encodeSymbol(uint32_t value, int k1, int k2, int k3)
137
{
138
139
140
141
142
  bool isZero = value == 0;
  arithmeticEncoder.encode(isZero, binaryModelIsZero[k1]);
  if (isZero) {
    return;
  }
143
144

  bool isOne = value == 1;
145
  arithmeticEncoder.encode(isOne, binaryModelIsOne[k2]);
146
147
148
149
150
  if (isOne) {
    return;
  }
  value -= 2;

151
  if (value < kAttributeResidualAlphabetSize) {
152
    symbolCoder[k3].encode(value, &arithmeticEncoder);
153
  } else {
154
    int alphabetSize = kAttributeResidualAlphabetSize;
155
    symbolCoder[k3].encode(alphabetSize, &arithmeticEncoder);
156
    arithmeticEncoder.encodeExpGolomb(
157
      value - alphabetSize, 0, binaryModel0, binaryModelDiff[k1]);
158
159
160
  }
}

161
162
163
//----------------------------------------------------------------------------

void
164
PCCResidualsEncoder::encode(uint32_t value0, uint32_t value1, uint32_t value2)
165
{
166
167
168
169
170
171
  if (value0 == value1 && value0 == value2) {
    value0--;
    value1--;
    value2--;
  }

172
173
174
175
176
177
178
  int b0 = (value0 == 0);
  int b1 = (value0 <= 1);
  int b2 = (value1 == 0);
  int b3 = (value1 <= 1);
  encodeSymbol(value0, 0, 0, 0);
  encodeSymbol(value1, 1 + b0, 1 + b1, 1);
  encodeSymbol(value2, 3 + (b0 << 1) + b2, 3 + (b1 << 1) + b3, 1);
179
180
181
182
183
184
185
}

//----------------------------------------------------------------------------

void
PCCResidualsEncoder::encode(uint32_t value)
{
186
  encodeSymbol(value - 1, 0, 0, 0);
187
188
}

189
//============================================================================
190
191
192
// An encapsulation of the entropy coding methods used in attribute coding

struct PCCResidualsEntropyEstimator {
193
194
  size_t freq0[kAttributeResidualAlphabetSize + 1];
  size_t freq1[kAttributeResidualAlphabetSize + 1];
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
  size_t symbolCount0;
  size_t symbolCount1;
  size_t isZero0Count;
  size_t isZero1Count;
  PCCResidualsEntropyEstimator() { init(); }
  void init();
  double bitsDetail(
    const uint32_t detail,
    const size_t symbolCount,
    const size_t* const freq) const;
  double bits(const uint32_t value0) const;
  void update(const uint32_t value0);
  double bits(
    const uint32_t value0, const uint32_t value1, const uint32_t value2) const;
  void
  update(const uint32_t value0, const uint32_t value1, const uint32_t value2);
};

//----------------------------------------------------------------------------
214

215
void
216
PCCResidualsEntropyEstimator::init()
217
{
218
  for (size_t i = 0; i <= kAttributeResidualAlphabetSize; ++i) {
219
220
221
    freq0[i] = 1;
    freq1[i] = 1;
  }
222
223
  symbolCount0 = kAttributeResidualAlphabetSize + 1;
  symbolCount1 = kAttributeResidualAlphabetSize + 1;
224
225
  isZero1Count = isZero0Count = symbolCount0 / 2;
}
226

227
228
229
230
231
232
233
234
235
//----------------------------------------------------------------------------

double
PCCResidualsEntropyEstimator::bitsDetail(
  const uint32_t detail,
  const size_t symbolCount,
  const size_t* const freq) const
{
  const uint32_t detailClipped =
236
    std::min(detail, uint32_t(kAttributeResidualAlphabetSize));
237
238
239
  const double pDetail =
    PCCClip(double(freq[detailClipped]) / symbolCount, 0.001, 0.999);
  double bits = -log2(pDetail);
240
241
  if (detail >= kAttributeResidualAlphabetSize) {
    const double x = double(detail) - double(kAttributeResidualAlphabetSize);
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
    bits += 2.0 * std::floor(log2(x + 1.0)) + 1.0;
  }
  return bits;
}

//----------------------------------------------------------------------------

double
PCCResidualsEntropyEstimator::bits(const uint32_t value0) const
{
  const bool isZero0 = value0 == 0;
  const double pIsZero0 = isZero0
    ? double(isZero0Count) / symbolCount0
    : double(symbolCount0 - isZero0Count) / symbolCount0;
  double bits = -log2(PCCClip(pIsZero0, 0.001, 0.999));
  if (!isZero0) {
    bits += bitsDetail(value0 - 1, symbolCount0, freq0);
259
  }
260
  return bits;
261
262
263
264
}

//----------------------------------------------------------------------------

265
266
267
268
269
270
void
PCCResidualsEntropyEstimator::update(const uint32_t value0)
{
  const bool isZero0 = value0 == 0;
  ++symbolCount0;
  if (!isZero0) {
271
    ++freq0[std::min(value0 - 1, uint32_t(kAttributeResidualAlphabetSize))];
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
  } else {
    ++isZero0Count;
  }
}

//----------------------------------------------------------------------------

double
PCCResidualsEntropyEstimator::bits(
  const uint32_t value0, const uint32_t value1, const uint32_t value2) const
{
  const bool isZero0 = value0 == 0;
  const double pIsZero0 = isZero0
    ? double(isZero0Count) / symbolCount0
    : double(symbolCount0 - isZero0Count) / symbolCount0;
  double bits = -log2(PCCClip(pIsZero0, 0.001, 0.999));
  if (!isZero0) {
    bits += bitsDetail(value0 - 1, symbolCount0, freq0);
  }

  const bool isZero1 = value1 == 0 && value2 == 0;
  const double pIsZero1 = isZero1
    ? double(isZero1Count) / symbolCount0
    : double(symbolCount0 - isZero1Count) / symbolCount0;
  bits -= log2(PCCClip(pIsZero1, 0.001, 0.999));
  if (!isZero1) {
    bits += bitsDetail(value1, symbolCount1, freq1);
    bits += bitsDetail(value2, symbolCount1, freq1);
  }
  return bits;
}

//----------------------------------------------------------------------------

void
PCCResidualsEntropyEstimator::update(
  const uint32_t value0, const uint32_t value1, const uint32_t value2)
{
  const bool isZero0 = value0 == 0;
  ++symbolCount0;
  if (!isZero0) {
313
    ++freq0[std::min(value0 - 1, uint32_t(kAttributeResidualAlphabetSize))];
314
315
316
317
318
319
320
  } else {
    ++isZero0Count;
  }

  const bool isZero1 = value1 == 0 && value2 == 0;
  symbolCount1 += 2;
  if (!isZero1) {
321
322
    ++freq1[std::min(value1, uint32_t(kAttributeResidualAlphabetSize))];
    ++freq1[std::min(value2, uint32_t(kAttributeResidualAlphabetSize))];
323
324
325
326
327
328
329
330
  } else {
    ++isZero1Count;
  }
}

//============================================================================
// AttributeEncoder Members

331
void
332
AttributeEncoder::encode(
333
  const SequenceParameterSet& sps,
334
  const AttributeDescription& desc,
335
  const AttributeParameterSet& attr_aps,
336
  const AttributeBrickHeader& abh,
337
  PCCPointSet3& pointCloud,
338
  PayloadBuffer* payload)
339
{
340
  std::vector<Quantizers> quantLayers = deriveQuantizerLayers(attr_aps, abh);
341

342
  PCCResidualsEncoder encoder;
343
  encoder.start(sps, int(pointCloud.getPointCount()));
344

345
  if (desc.attr_num_dimensions == 1) {
346
347
    switch (attr_aps.attr_encoding) {
    case AttributeEncoding::kRAHTransform:
348
      encodeReflectancesTransformRaht(
349
        desc, attr_aps, quantLayers, pointCloud, encoder);
350
      break;
351

352
    case AttributeEncoding::kPredictingTransform:
353
      encodeReflectancesPred(desc, attr_aps, quantLayers, pointCloud, encoder);
354
      break;
355

356
    case AttributeEncoding::kLiftingTransform:
357
      encodeReflectancesLift(desc, attr_aps, quantLayers, pointCloud, encoder);
358
359
      break;
    }
360
  } else if (desc.attr_num_dimensions == 3) {
361
362
    switch (attr_aps.attr_encoding) {
    case AttributeEncoding::kRAHTransform:
363
364
      encodeColorsTransformRaht(
        desc, attr_aps, quantLayers, pointCloud, encoder);
365
366
367
      break;

    case AttributeEncoding::kPredictingTransform:
368
      encodeColorsPred(desc, attr_aps, quantLayers, pointCloud, encoder);
369
370
371
      break;

    case AttributeEncoding::kLiftingTransform:
372
      encodeColorsLift(desc, attr_aps, quantLayers, pointCloud, encoder);
373
374
375
      break;
    }
  } else {
376
    assert(desc.attr_num_dimensions == 1 || desc.attr_num_dimensions == 3);
377
  }
378

379
380
381
382
  uint32_t acDataLen = encoder.stop();
  std::copy_n(
    encoder.arithmeticEncoder.buffer(), acDataLen,
    std::back_inserter(*payload));
383
384
385
386
}

//----------------------------------------------------------------------------

387
388
389
390
int64_t
AttributeEncoder::computeReflectanceResidual(
  const uint64_t reflectance,
  const uint64_t predictedReflectance,
391
  const Quantizer& quant)
392
393
394
{
  const int64_t quantAttValue = reflectance;
  const int64_t quantPredAttValue = predictedReflectance;
395
396
  const int64_t delta = quant.quantize(
    (quantAttValue - quantPredAttValue) << kFixedPointAttributeShift);
397

398
  return IntToUInt(delta);
399
400
401
402
403
404
}

//----------------------------------------------------------------------------

void
AttributeEncoder::computeReflectancePredictionWeights(
405
  const AttributeParameterSet& aps,
406
  const PCCPointSet3& pointCloud,
407
408
  const std::vector<uint32_t>& indexesLOD,
  const uint32_t predictorIndex,
409
410
  PCCPredictor& predictor,
  PCCResidualsEncoder& encoder,
411
  PCCResidualsEntropyEstimator& context,
412
  const Quantizer& quant)
413
414
{
  predictor.computeWeights();
415
  predictor.maxDiff = 0;
416
417
418
  if (predictor.neighborCount > 1) {
    int64_t minValue = 0;
    int64_t maxValue = 0;
419
    for (size_t i = 0; i < predictor.neighborCount; ++i) {
420
      const uint64_t reflectanceNeighbor = pointCloud.getReflectance(
421
        indexesLOD[predictor.neighbors[i].predictorIndex]);
422
423
424
425
426
427
428
429
      if (i == 0 || reflectanceNeighbor < minValue) {
        minValue = reflectanceNeighbor;
      }
      if (i == 0 || reflectanceNeighbor > maxValue) {
        maxValue = reflectanceNeighbor;
      }
    }
    const int64_t maxDiff = maxValue - minValue;
430
    predictor.maxDiff = maxDiff;
431
    if (maxDiff >= aps.adaptive_prediction_threshold) {
432
      uint64_t attrValue =
433
        pointCloud.getReflectance(indexesLOD[predictorIndex]);
434
435
436

      // base case: weighted average of n neighbours
      predictor.predMode = 0;
437
      uint64_t attrPred = predictor.predictReflectance(pointCloud, indexesLOD);
438
      int64_t attrResidualQuant =
439
        computeReflectanceResidual(attrValue, attrPred, quant);
440

441
      double best_score = attrResidualQuant
442
        + kAttrPredLambdaR * (quant.stepSize() >> kFixedPointAttributeShift);
443
444
445
446
447

      for (int i = 0; i < predictor.neighborCount; i++) {
        if (i == aps.max_num_direct_predictors)
          break;

448
449
        attrPred = pointCloud.getReflectance(
          indexesLOD[predictor.neighbors[i].predictorIndex]);
450
        attrResidualQuant =
451
          computeReflectanceResidual(attrValue, attrPred, quant);
452
453

        double idxBits = i + (i == aps.max_num_direct_predictors - 1 ? 1 : 2);
454
        double score = attrResidualQuant
455
456
          + idxBits * kAttrPredLambdaR
            * (quant.stepSize() >> kFixedPointAttributeShift);
457
458
459
460
461
462
463

        if (score < best_score) {
          best_score = score;
          predictor.predMode = i + 1;
          // NB: setting predictor.neighborCount = 1 will cause issues
          // with reconstruction.
        }
464
465
466
467
468
469
470
      }
    }
  }
}

//----------------------------------------------------------------------------

471
void
472
AttributeEncoder::encodeReflectancesPred(
473
  const AttributeDescription& desc,
474
  const AttributeParameterSet& aps,
475
  const std::vector<Quantizers>& quantLayers,
476
477
478
  PCCPointSet3& pointCloud,
  PCCResidualsEncoder& encoder)
{
479
480
  const uint32_t pointCount = pointCloud.getPointCount();
  std::vector<PCCPredictor> predictors;
481
482
  std::vector<uint32_t> numberOfPointsPerLOD;
  std::vector<uint32_t> indexesLOD;
483

484
  buildPredictorsFast(
485
    aps, pointCloud, 0, predictors, numberOfPointsPerLOD, indexesLOD);
486

487
  const int64_t clipMax = (1ll << desc.attr_bitdepth) - 1;
488
  PCCResidualsEntropyEstimator context;
489
490
491
492
493
494
  int zero_cnt = 0;
  std::vector<int> zerorun;
  zerorun.reserve(pointCount);
  std::vector<uint32_t> residual;
  residual.resize(pointCount);

495
  int quantLayer = 0;
496
497
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
498
499
500
501
    if (predictorIndex == numberOfPointsPerLOD[quantLayer]) {
      quantLayer = std::min(int(quantLayers.size()) - 1, quantLayer + 1);
    }
    auto& quant = quantLayers[quantLayer];
502
    auto& predictor = predictors[predictorIndex];
503

504
    computeReflectancePredictionWeights(
505
      aps, pointCloud, indexesLOD, predictorIndex, predictor, encoder, context,
506
      quant[0]);
507

508
    const uint32_t pointIndex = indexesLOD[predictorIndex];
509
    const uint64_t reflectance = pointCloud.getReflectance(pointIndex);
510
    const attr_t predictedReflectance =
511
      predictor.predictReflectance(pointCloud, indexesLOD);
512
513
    const int64_t quantAttValue = reflectance;
    const int64_t quantPredAttValue = predictedReflectance;
514
515
    const int64_t delta = quant[0].quantize(
      (quantAttValue - quantPredAttValue) << kFixedPointAttributeShift);
516
    const uint32_t attValue0 = uint32_t(IntToUInt(long(delta)));
517
518
    const int64_t reconstructedDelta =
      divExp2RoundHalfUp(quant[0].scale(delta), kFixedPointAttributeShift);
519
520
    const int64_t reconstructedQuantAttValue =
      quantPredAttValue + reconstructedDelta;
521
522
    const attr_t reconstructedReflectance =
      attr_t(PCCClip(reconstructedQuantAttValue, int64_t(0), clipMax));
523

524
525
526
527
528
529
530
    if (!attValue0)
      ++zero_cnt;
    else {
      zerorun.push_back(zero_cnt);
      zero_cnt = 0;
    }
    residual[predictorIndex] = attValue0;
531
    pointCloud.setReflectance(pointIndex, reconstructedReflectance);
532
  }
533
534
535
536
537
538
539
540
541

  zerorun.push_back(zero_cnt);
  int run_index = 0;
  encoder.encodeZeroCnt(zerorun[run_index], pointCount);
  zero_cnt = zerorun[run_index++];

  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
    auto& predictor = predictors[predictorIndex];
542
    if (predictor.maxDiff >= aps.adaptive_prediction_threshold) {
543
544
545
546
547
548
549
550
551
552
553
554
      encoder.encodePredMode(
        predictor.predMode, aps.max_num_direct_predictors);
    }
    if (zero_cnt > 0)
      zero_cnt--;
    else {
      encoder.encode(residual[predictorIndex]);
      if (predictorIndex != pointCount - 1)
        encoder.encodeZeroCnt(zerorun[run_index], pointCount);
      zero_cnt = zerorun[run_index++];
    }
  }
555
556
557
558
}

//----------------------------------------------------------------------------

559
Vec3<int64_t>
560
AttributeEncoder::computeColorResiduals(
561
562
  const Vec3<attr_t> color,
  const Vec3<attr_t> predictedColor,
563
  const Quantizers& quant)
564
{
565
  Vec3<int64_t> residuals;
566
567
  const int64_t quantAttValue = color[0];
  const int64_t quantPredAttValue = predictedColor[0];
568
569
  const int64_t delta = quant[0].quantize(
    (quantAttValue - quantPredAttValue) << kFixedPointAttributeShift);
570
  residuals[0] = IntToUInt(delta);
571
572
573
  for (size_t k = 1; k < 3; ++k) {
    const int64_t quantAttValue = color[k];
    const int64_t quantPredAttValue = predictedColor[k];
574
575
    const int64_t delta = quant[1].quantize(
      (quantAttValue - quantPredAttValue) << kFixedPointAttributeShift);
576
    residuals[k] = IntToUInt(delta);
577
578
579
580
581
582
583
584
  }
  return residuals;
}

//----------------------------------------------------------------------------

void
AttributeEncoder::computeColorPredictionWeights(
585
  const AttributeParameterSet& aps,
586
  const PCCPointSet3& pointCloud,
587
588
  const std::vector<uint32_t>& indexesLOD,
  const uint32_t predictorIndex,
589
590
  PCCPredictor& predictor,
  PCCResidualsEncoder& encoder,
591
  PCCResidualsEntropyEstimator& context,
592
  const Quantizers& quant)
593
594
{
  predictor.computeWeights();
595
  predictor.maxDiff = 0;
596
597
598
  if (predictor.neighborCount > 1) {
    int64_t minValue[3] = {0, 0, 0};
    int64_t maxValue[3] = {0, 0, 0};
599
    for (int i = 0; i < predictor.neighborCount; ++i) {
600
      const Vec3<attr_t> colorNeighbor =
601
        pointCloud.getColor(indexesLOD[predictor.neighbors[i].predictorIndex]);
602
603
604
605
606
607
608
609
610
611
612
613
      for (size_t k = 0; k < 3; ++k) {
        if (i == 0 || colorNeighbor[k] < minValue[k]) {
          minValue[k] = colorNeighbor[k];
        }
        if (i == 0 || colorNeighbor[k] > maxValue[k]) {
          maxValue[k] = colorNeighbor[k];
        }
      }
    }
    const int64_t maxDiff = (std::max)(
      maxValue[2] - minValue[2],
      (std::max)(maxValue[0] - minValue[0], maxValue[1] - minValue[1]));
614
615
    predictor.maxDiff = maxDiff;

616
    if (maxDiff >= aps.adaptive_prediction_threshold) {
617
      Vec3<attr_t> attrValue = pointCloud.getColor(indexesLOD[predictorIndex]);
618
619
620

      // base case: weighted average of n neighbours
      predictor.predMode = 0;
621
      Vec3<attr_t> attrPred = predictor.predictColor(pointCloud, indexesLOD);
622
      Vec3<int64_t> attrResidualQuant =
623
        computeColorResiduals(attrValue, attrPred, quant);
624
625

      double best_score = attrResidualQuant[0] + attrResidualQuant[1]
626
        + attrResidualQuant[2]
627
628
        + kAttrPredLambdaC
          * (double)(quant[0].stepSize() >> kFixedPointAttributeShift);
629
630
631
632
633

      for (int i = 0; i < predictor.neighborCount; i++) {
        if (i == aps.max_num_direct_predictors)
          break;

634
635
        attrPred = pointCloud.getColor(
          indexesLOD[predictor.neighbors[i].predictorIndex]);
636
        attrResidualQuant = computeColorResiduals(attrValue, attrPred, quant);
637
638
639

        double idxBits = i + (i == aps.max_num_direct_predictors - 1 ? 1 : 2);
        double score = attrResidualQuant[0] + attrResidualQuant[1]
640
          + attrResidualQuant[2]
641
642
          + idxBits * kAttrPredLambdaC
            * (quant[0].stepSize() >> kFixedPointAttributeShift);
643
644
645
646
647
648
649

        if (score < best_score) {
          best_score = score;
          predictor.predMode = i + 1;
          // NB: setting predictor.neighborCount = 1 will cause issues
          // with reconstruction.
        }
650
651
652
653
654
655
656
      }
    }
  }
}

//----------------------------------------------------------------------------

657
void
658
AttributeEncoder::encodeColorsPred(
659
  const AttributeDescription& desc,
660
  const AttributeParameterSet& aps,
661
  const std::vector<Quantizers>& quantLayers,
662
663
664
  PCCPointSet3& pointCloud,
  PCCResidualsEncoder& encoder)
{
665
  const size_t pointCount = pointCloud.getPointCount();
666
  std::vector<PCCPredictor> predictors;
667
668
  std::vector<uint32_t> numberOfPointsPerLOD;
  std::vector<uint32_t> indexesLOD;
669

670
  buildPredictorsFast(
671
    aps, pointCloud, 0, predictors, numberOfPointsPerLOD, indexesLOD);
672

673
674
675
676
  Vec3<int64_t> clipMax{(1 << desc.attr_bitdepth) - 1,
                        (1 << desc.attr_bitdepth_secondary) - 1,
                        (1 << desc.attr_bitdepth_secondary) - 1};

677
678
  uint32_t values[3];
  PCCResidualsEntropyEstimator context;
679
680
681
682
683
684
685
  int zero_cnt = 0;
  std::vector<int> zerorun;
  std::vector<uint32_t> residual[3];
  for (int i = 0; i < 3; i++) {
    residual[i].resize(pointCount);
  }

686
  int quantLayer = 0;
687
688
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
689
690
691
692
    if (predictorIndex == numberOfPointsPerLOD[quantLayer]) {
      quantLayer = std::min(int(quantLayers.size()) - 1, quantLayer + 1);
    }
    auto& quant = quantLayers[quantLayer];
693
    auto& predictor = predictors[predictorIndex];
694

695
    computeColorPredictionWeights(
696
      aps, pointCloud, indexesLOD, predictorIndex, predictor, encoder, context,
697
      quant);
698
    const auto pointIndex = indexesLOD[predictorIndex];
699
700
    const Vec3<attr_t> color = pointCloud.getColor(pointIndex);
    const Vec3<attr_t> predictedColor =
701
      predictor.predictColor(pointCloud, indexesLOD);
702

703
    Vec3<attr_t> reconstructedColor;
704
    int64_t residual0 = 0;
705
706
707
708
709
710
711
712
    for (int k = 0; k < 3; ++k) {
      const auto& q = quant[std::min(k, 1)];
      int64_t residual = color[k] - predictedColor[k];

      int64_t residualQ = q.quantize(residual << kFixedPointAttributeShift);
      int64_t residualR =
        divExp2RoundHalfUp(q.scale(residualQ), kFixedPointAttributeShift);

713
714
715
716
717
718
719
720
721
722
      if (aps.inter_component_prediction_enabled_flag && k > 0) {
        residual = residualR - residual0;
        residualQ = q.quantize(residual << kFixedPointAttributeShift);
        residualR = residual0
          + divExp2RoundHalfUp(q.scale(residualQ), kFixedPointAttributeShift);
      }

      if (k == 0)
        residual0 = residualR;

723
724
725
      values[k] = uint32_t(IntToUInt(long(residualQ)));

      int64_t recon = predictedColor[k] + residualR;
726
      reconstructedColor[k] = attr_t(PCCClip(recon, int64_t(0), clipMax[k]));
727
    }
728
    pointCloud.setColor(pointIndex, reconstructedColor);
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748

    if (!values[0] && !values[1] && !values[2]) {
      ++zero_cnt;
    } else {
      zerorun.push_back(zero_cnt);
      zero_cnt = 0;
    }

    for (int i = 0; i < 3; i++) {
      residual[i][predictorIndex] = values[i];
    }
  }

  zerorun.push_back(zero_cnt);
  int run_index = 0;
  encoder.encodeZeroCnt(zerorun[run_index], pointCount);
  zero_cnt = zerorun[run_index++];
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
    auto& predictor = predictors[predictorIndex];
749
    if (predictor.maxDiff >= aps.adaptive_prediction_threshold) {
750
751
752
753
754
755
756
757
758
759
760
761
762
763
      encoder.encodePredMode(
        predictor.predMode, aps.max_num_direct_predictors);
    }
    if (zero_cnt > 0)
      zero_cnt--;
    else {
      for (size_t k = 0; k < 3; k++)
        values[k] = residual[k][predictorIndex];

      encoder.encode(values[0], values[1], values[2]);
      if (predictorIndex != pointCount - 1)
        encoder.encodeZeroCnt(zerorun[run_index], pointCount);
      zero_cnt = zerorun[run_index++];
    }
764
765
766
  }
}

767
768
//----------------------------------------------------------------------------

769
770
void
AttributeEncoder::encodeReflectancesTransformRaht(
771
  const AttributeDescription& desc,
772
  const AttributeParameterSet& aps,
773
  const std::vector<Quantizers>& quantLayers,
774
775
776
  PCCPointSet3& pointCloud,
  PCCResidualsEncoder& encoder)
{
777
778
779
  const int voxelCount = int(pointCloud.getPointCount());
  std::vector<MortonCodeWithIndex> packedVoxel(voxelCount);
  for (int n = 0; n < voxelCount; n++) {
780
    packedVoxel[n].mortonCode = mortonAddr(pointCloud[n], 0);
781
782
783
784
785
    packedVoxel[n].index = n;
  }
  sort(packedVoxel.begin(), packedVoxel.end());

  // Allocate arrays.
786
  int64_t* mortonCode = new int64_t[voxelCount];
787
  const int attribCount = 1;
David Flynn's avatar
David Flynn committed
788
789
  int* attributes = new int[attribCount * voxelCount];
  int* coefficients = new int[attribCount * voxelCount];
790
791
792
793

  // Populate input arrays.
  for (int n = 0; n < voxelCount; n++) {
    mortonCode[n] = packedVoxel[n].mortonCode;
794
795
    const auto reflectance = pointCloud.getReflectance(packedVoxel[n].index);
    attributes[attribCount * n] = reflectance;
796
797
798
799
  }

  // Transform.
  regionAdaptiveHierarchicalTransform(
800
    aps.raht_prediction_enabled_flag, quantLayers, mortonCode, attributes,
801
    attribCount, voxelCount, coefficients);
802

803
  // Entropy encode.
804
  int zero_cnt = 0;
805
  uint32_t value;
806
  for (int n = 0; n < voxelCount; ++n) {
David Flynn's avatar
David Flynn committed
807
    const int64_t detail = IntToUInt(coefficients[n]);
808
    assert(detail < std::numeric_limits<uint32_t>::max());
809
    value = uint32_t(detail);
810
811
812
813
814
815
816
    if (!value)
      ++zero_cnt;
    else {
      encoder.encodeZeroCnt(zero_cnt, voxelCount);
      encoder.encode(value);
      zero_cnt = 0;
    }
817
  }
818
  encoder.encodeZeroCnt(zero_cnt, voxelCount);
819

820
821
  const int64_t maxReflectance = (1 << desc.attr_bitdepth) - 1;
  const int64_t minReflectance = 0;
822
  for (int n = 0; n < voxelCount; n++) {
David Flynn's avatar
David Flynn committed
823
    int64_t val = attributes[attribCount * n];
824
825
    const attr_t reflectance =
      attr_t(PCCClip(val, minReflectance, maxReflectance));
826
    pointCloud.setReflectance(packedVoxel[n].index, reflectance);
827
828
829
830
831
  }

  // De-allocate arrays.
  delete[] mortonCode;
  delete[] attributes;
David Flynn's avatar
David Flynn committed
832
  delete[] coefficients;
833
834
835
836
}

//----------------------------------------------------------------------------

837
838
void
AttributeEncoder::encodeColorsTransformRaht(
839
  const AttributeDescription& desc,
840
  const AttributeParameterSet& aps,
841
  const std::vector<Quantizers>& quantLayers,
842
843
844
  PCCPointSet3& pointCloud,
  PCCResidualsEncoder& encoder)
{
845
846
847
  const int voxelCount = int(pointCloud.getPointCount());
  std::vector<MortonCodeWithIndex> packedVoxel(voxelCount);
  for (int n = 0; n < voxelCount; n++) {
848
    packedVoxel[n].mortonCode = mortonAddr(pointCloud[n], 0);
849
850
851
852
853
    packedVoxel[n].index = n;
  }
  sort(packedVoxel.begin(), packedVoxel.end());

  // Allocate arrays.
854
  int64_t* mortonCode = new int64_t[voxelCount];
855
  const int attribCount = 3;
David Flynn's avatar
David Flynn committed
856
857
  int* attributes = new int[attribCount * voxelCount];
  int* coefficients = new int[attribCount * voxelCount];
858
859
860
861
862
863
864
865
866
867
868
869

  // Populate input arrays.
  for (int n = 0; n < voxelCount; n++) {
    mortonCode[n] = packedVoxel[n].mortonCode;
    const auto color = pointCloud.getColor(packedVoxel[n].index);
    attributes[attribCount * n] = color[0];
    attributes[attribCount * n + 1] = color[1];
    attributes[attribCount * n + 2] = color[2];
  }

  // Transform.
  regionAdaptiveHierarchicalTransform(
870
    aps.raht_prediction_enabled_flag, quantLayers, mortonCode, attributes,
871
    attribCount, voxelCount, coefficients);
872
873

  // Entropy encode.
874
  uint32_t values[attribCount];
875
  int zero_cnt = 0;
876
  for (int n = 0; n < voxelCount; ++n) {
877
    for (int d = 0; d < attribCount; ++d) {
David Flynn's avatar
David Flynn committed
878
      const int64_t detail = IntToUInt(coefficients[voxelCount * d + n]);
879
880
      assert(detail < std::numeric_limits<uint32_t>::max());
      values[d] = uint32_t(detail);
881
    }
882
883
884
885
886
887
888
    if (!values[0] && !values[1] && !values[2])
      ++zero_cnt;
    else {
      encoder.encodeZeroCnt(zero_cnt, voxelCount);
      encoder.encode(values[0], values[1], values[2]);
      zero_cnt = 0;
    }
889
  }
890
891
  encoder.encodeZeroCnt(zero_cnt, voxelCount);

892
893
894
895
  Vec3<int> clipMax{(1 << desc.attr_bitdepth) - 1,
                    (1 << desc.attr_bitdepth_secondary) - 1,
                    (1 << desc.attr_bitdepth_secondary) - 1};

896
  for (int n = 0; n < voxelCount; n++) {
David Flynn's avatar
David Flynn committed
897
898
899
    const int r = attributes[attribCount * n];
    const int g = attributes[attribCount * n + 1];
    const int b = attributes[attribCount * n + 2];
900
    Vec3<attr_t> color;
901
902
903
    color[0] = attr_t(PCCClip(r, 0, clipMax[0]));
    color[1] = attr_t(PCCClip(g, 0, clipMax[1]));
    color[2] = attr_t(PCCClip(b, 0, clipMax[2]));
904
905
906
907
908
909
    pointCloud.setColor(packedVoxel[n].index, color);
  }

  // De-allocate arrays.
  delete[] mortonCode;
  delete[] attributes;
David Flynn's avatar
David Flynn committed
910
  delete[] coefficients;
911
}
912
913
914
915
//----------------------------------------------------------------------------

void
AttributeEncoder::encodeColorsLift(
916
  const AttributeDescription& desc,
917
  const AttributeParameterSet& aps,
918
  const std::vector<Quantizers>& quantLayers,
919
920
921
922
  PCCPointSet3& pointCloud,
  PCCResidualsEncoder& encoder)
{
  const size_t pointCount = pointCloud.getPointCount();
923
  std::vector<PCCPredictor> predictors;
924
925
  std::vector<uint32_t> numberOfPointsPerLOD;
  std::vector<uint32_t> indexesLOD;
926

927
  buildPredictorsFast(
928
    aps, pointCloud, 0, predictors, numberOfPointsPerLOD, indexesLOD);
929

930
931
932
933
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
    predictors[predictorIndex].computeWeights();
  }
934
  std::vector<uint64_t> weights;
935
936
937
938
939
940
  if (!aps.scalable_lifting_enabled_flag) {
    PCCComputeQuantizationWeights(predictors, weights);
  } else {
    computeQuantizationWeightsScalable(
      predictors, numberOfPointsPerLOD, pointCount, 0, weights);
  }
941
  const size_t lodCount = numberOfPointsPerLOD.size();
942
  std::vector<Vec3<int64_t>> colors;
943
944
945
946
947
  colors.resize(pointCount);

  for (size_t index = 0; index < pointCount; ++index) {
    const auto& color = pointCloud.getColor(indexesLOD[index]);
    for (size_t d = 0; d < 3; ++d) {
948
      colors[index][d] = int32_t(color[d]) << kFixedPointAttributeShift;
949
950
951
952
953
954
955
956
957
958
959
960
    }
  }

  for (size_t i = 0; (i + 1) < lodCount; ++i) {
    const size_t lodIndex = lodCount - i - 1;
    const size_t startIndex = numberOfPointsPerLOD[lodIndex - 1];
    const size_t endIndex = numberOfPointsPerLOD[lodIndex];
    PCCLiftPredict(predictors, startIndex, endIndex, true, colors);
    PCCLiftUpdate(predictors, weights, startIndex, endIndex, true, colors);
  }

  // compress
961
  int zero_cnt = 0;
962
  int quantLayer = 0;
963
964
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
965
966
967
968
969
    if (predictorIndex == numberOfPointsPerLOD[quantLayer]) {
      quantLayer = std::min(int(quantLayers.size()) - 1, quantLayer + 1);
    }
    auto& quant = quantLayers[quantLayer];

970
    const int64_t quantWeight = weights[predictorIndex];
971
    auto& color = colors[predictorIndex];
972
    const int64_t delta = quant[0].quantize(color[0] * quantWeight);
973
    const int64_t detail = IntToUInt(delta);
974
    assert(detail < std::numeric_limits<uint32_t>::max());
975
    const int64_t reconstructedDelta = quant[0].scale(delta);
976
977
978
979
    color[0] = reconstructedDelta / quantWeight;
    uint32_t values[3];
    values[0] = uint32_t(detail);
    for (size_t d = 1; d < 3; ++d) {
980
      const int64_t delta = quant[1].quantize(color[d] * quantWeight);
981
      const int64_t detail = IntToUInt(delta);
982
      assert(detail < std::numeric_limits<uint32_t>::max());
983
      const int64_t reconstructedDelta = quant[1].scale(delta);
984
985
986
      color[d] = reconstructedDelta / quantWeight;
      values[d] = uint32_t(detail);
    }
987
988
989
990
991
992
993
    if (!values[0] && !values[1] && !values[2])
      ++zero_cnt;
    else {
      encoder.encodeZeroCnt(zero_cnt, pointCount);
      encoder.encode(values[0], values[1], values[2]);
      zero_cnt = 0;
    }
994
  }
995
  encoder.encodeZeroCnt(zero_cnt, pointCount);
996
997
998
999
1000
1001
1002
1003

  // reconstruct
  for (size_t lodIndex = 1; lodIndex < lodCount; ++lodIndex) {
    const size_t startIndex = numberOfPointsPerLOD[lodIndex - 1];
    const size_t endIndex = numberOfPointsPerLOD[lodIndex];
    PCCLiftUpdate(predictors, weights, startIndex, endIndex, false, colors);
    PCCLiftPredict(predictors, startIndex, endIndex, false, colors);
  }
1004