AttributeEncoder.cpp 35.9 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 uint16_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
523
    const uint16_t reconstructedReflectance =
      uint16_t(PCCClip(reconstructedQuantAttValue, int64_t(0), clipMax));

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<uint8_t> color,
  const Vec3<uint8_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<uint8_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
618
      Vec3<uint8_t> attrValue =
        pointCloud.getColor(indexesLOD[predictorIndex]);
619
620
621

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

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

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

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

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

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

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

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

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

674
  const int64_t clipMax = (1ll << desc.attr_bitdepth) - 1;
675
676
  uint32_t values[3];
  PCCResidualsEntropyEstimator context;
677
678
679
680
681
682
683
  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);
  }

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

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

701
    Vec3<uint8_t> reconstructedColor;
702
    int64_t residual0 = 0;
703
704
705
706
707
708
709
710
    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);

711
712
713
714
715
716
717
718
719
720
      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;

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

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

    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];
747
    if (predictor.maxDiff >= aps.adaptive_prediction_threshold) {
748
749
750
751
752
753
754
755
756
757
758
759
760
761
      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++];
    }
762
763
764
  }
}

765
766
//----------------------------------------------------------------------------

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

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

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

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

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

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

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

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

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

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

  // 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(
868
    aps.raht_prediction_enabled_flag, quantLayers, mortonCode, attributes,
869
    attribCount, voxelCount, coefficients);
870
871

  // Entropy encode.
872
  uint32_t values[attribCount];
873
  int zero_cnt = 0;
874
  for (int n = 0; n < voxelCount; ++n) {
875
    for (int d = 0; d < attribCount; ++d) {
David Flynn's avatar
David Flynn committed
876
      const int64_t detail = IntToUInt(coefficients[voxelCount * d + n]);
877
878
      assert(detail < std::numeric_limits<uint32_t>::max());
      values[d] = uint32_t(detail);
879
    }
880
881
882
883
884
885
886
    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;
    }
887
  }
888
889
  encoder.encodeZeroCnt(zero_cnt, voxelCount);

890
  const int clipMax = (1 << desc.attr_bitdepth) - 1;
891
  for (int n = 0; n < voxelCount; n++) {
David Flynn's avatar
David Flynn committed
892
893
894
    const int r = attributes[attribCount * n];
    const int g = attributes[attribCount * n + 1];
    const int b = attributes[attribCount * n + 2];
895
    Vec3<uint8_t> color;
896
897
898
    color[0] = uint8_t(PCCClip(r, 0, clipMax));
    color[1] = uint8_t(PCCClip(g, 0, clipMax));
    color[2] = uint8_t(PCCClip(b, 0, clipMax));
899
900
901
902
903
904
    pointCloud.setColor(packedVoxel[n].index, color);
  }

  // De-allocate arrays.
  delete[] mortonCode;
  delete[] attributes;
David Flynn's avatar
David Flynn committed
905
  delete[] coefficients;
906
}
907
908
909
910
//----------------------------------------------------------------------------

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

922
  buildPredictorsFast(
923
    aps, pointCloud, 0, predictors, numberOfPointsPerLOD, indexesLOD);
924

925
926
927
928
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
    predictors[predictorIndex].computeWeights();
  }
929
  std::vector<uint64_t> weights;
930
931
932
933
934
935
  if (!aps.scalable_lifting_enabled_flag) {
    PCCComputeQuantizationWeights(predictors, weights);
  } else {
    computeQuantizationWeightsScalable(
      predictors, numberOfPointsPerLOD, pointCount, 0, weights);
  }
936
  const size_t lodCount = numberOfPointsPerLOD.size();
937
  std::vector<Vec3<int64_t>> colors;
938
939
940
941
942
  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) {
943
      colors[index][d] = int32_t(color[d]) << kFixedPointAttributeShift;
944
945
946
947
948
949
950
951
952
953
954
955
    }
  }

  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
956
  int zero_cnt = 0;
957
  int quantLayer = 0;
958
959
  for (size_t predictorIndex = 0; predictorIndex < pointCount;
       ++predictorIndex) {
960
961
962
963
964
    if (predictorIndex == numberOfPointsPerLOD[quantLayer]) {
      quantLayer = std::min(int(quantLayers.size()) - 1, quantLayer + 1);
    }
    auto& quant = quantLayers[quantLayer];

965
    const int64_t quantWeight = weights[predictorIndex];
966
    auto& color = colors[predictorIndex];
967
    const int64_t delta = quant[0].quantize(color[0] * quantWeight);
968
    const int64_t detail = IntToUInt(delta);
969
    assert(detail < std::numeric_limits<uint32_t>::max());
970
    const int64_t reconstructedDelta = quant[0].scale(delta);
971
972
973
974
    color[0] = reconstructedDelta / quantWeight;
    uint32_t values[3];
    values[0] = uint32_t(detail);
    for (size_t d = 1; d < 3; ++d) {
975
      const int64_t delta = quant[1].quantize(color[d] * quantWeight);
976
      const int64_t detail = IntToUInt(delta);
977
      assert(detail < std::numeric_limits<uint32_t>::max());
978
      const int64_t reconstructedDelta = quant[1].scale(delta);
979
980
981
      color[d] = reconstructedDelta / quantWeight;
      values[d] = uint32_t(detail);
    }
982
983
984
985
986
987
988
    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;
    }
989
  }
990
  encoder.encodeZeroCnt(zero_cnt, pointCount);
991
992
993
994
995
996
997
998

  // 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);
  }
999

1000
  const int64_t clipMax = (1 << desc.attr_bitdepth) - 1;
1001
  for (size_t f = 0; f < pointCount; ++f) {
1002
1003
    const auto color0 =
      divExp2RoundHalfInf(colors[f], kFixedPointAttributeShift);