int main()

in vireo/tools/psnr/main.cpp [86:213]


int main(int argc, const char* argv[]) {
  __try {
    // Parse arguments
    if (argc < 3) {
      print_usage(common::Path::Filename(argv[0]));
      return 1;
    }
    Config config;
    if (parse_arguments(argc, argv, config)) {
      return 1;
    }

    // Read input files
    demux::Movie movie1(config.ref);
    demux::Movie movie2(config.test);

    const auto width1 = movie1.video_track.settings().width;
    const auto height1 = movie1.video_track.settings().height;
    const auto width2 = movie2.video_track.settings().width;
    const auto height2 = movie2.video_track.settings().height;
    const bool stretch = (width1 != width2) || (height1 != height2);

    // Decode frames and compare
    decode::Video decoder1(movie1.video_track);
    decode::Video decoder2(movie2.video_track);

    // Filter visible frames
    auto samples1 = movie1.video_track.filter([&edit_boxes = movie1.video_track.edit_boxes()](decode::Sample& sample){ return common::EditBox::Plays(edit_boxes, sample.pts); });
    auto samples2 = movie2.video_track.filter([&edit_boxes = movie2.video_track.edit_boxes()](decode::Sample& sample){ return common::EditBox::Plays(edit_boxes, sample.pts); });
    auto frames1 = decoder1.filter([&edit_boxes = movie1.video_track.edit_boxes()](frame::Frame& frm){ return common::EditBox::Plays(edit_boxes, frm.pts); });
    auto frames2 = decoder2.filter([&edit_boxes = movie2.video_track.edit_boxes()](frame::Frame& frm){ return common::EditBox::Plays(edit_boxes, frm.pts); });

    // Create a mapping between frames1 to frames2 according to PTS
    unordered_map<int, uint32_t> time_to_index2;
    vector<pair<uint32_t, uint32_t>> index1_to_index2;
    const int scale = 100;  // compare time up to two decimal points
    for (uint32_t index2 = 0; index2 < frames2.count(); ++index2) {
      const auto pts = common::EditBox::RealPts(movie2.video_track.edit_boxes(), frames2(index2).pts);
      const int time = (int)((float)pts * scale / frames2.settings().timescale);
      THROW_IF(time_to_index2.find(time) != time_to_index2.end(), Invalid);
      time_to_index2[time] = index2;
    }
    for (uint32_t index1 = 0; index1 < frames1.count(); ++index1) {
      const auto pts = common::EditBox::RealPts(movie1.video_track.edit_boxes(), frames1(index1).pts);
      const int time = (int)((float)pts * scale / frames1.settings().timescale);
      auto index2_search = time_to_index2.find(time);
      if (index2_search != time_to_index2.end()) {
        index1_to_index2.push_back(make_pair(index1, index2_search->second));
      }
    }

    // Check if there are frames that match between videos
    if (index1_to_index2.size() == 0) {
      cerr << "Videos do not contain any matching frames! Cannot compute PSNR" << endl;
      return 1;
    }
    cout << "Calculating PSNR over " << index1_to_index2.size() << " frames" << endl;
    if (index1_to_index2.size() < frames1.count()) {
      cout << "Warning: " << (frames1.count() - index1_to_index2.size()) << " frames of reference video not used" << endl;
    }
    if (index1_to_index2.size() < frames2.count()) {
      cout << "Warning: " << (frames2.count() - index1_to_index2.size()) << " frames of test video not used" << endl;
    }

    // Calculate PSNR
    double total_sse = 0.0;
    double total_num_pixels = 0.0;
    for (uint32_t i = 0; i < index1_to_index2.size(); ++i) {
      const uint32_t index1 = index1_to_index2[i].first;
      const uint32_t index2 = index1_to_index2[i].second;

      uint64_t sse_per_frame = 0;
      uint64_t num_pixels_per_frame = 0;
      if (samples1(index1).nal() != samples2(index2).nal()) {
        const auto& ref_frame = frames1(index1).yuv();
        const auto& test_frame = stretch ? frames2(index2).yuv().stretch(width1, width2, height1, height2) : frames2(index2).yuv();

        for (auto p: enumeration::Enum<frame::PlaneIndex>(frame::Y, frame::V)) {
          const auto& ref_plane = ref_frame.plane(p);
          const auto& test_plane = test_frame.plane(p);

          const auto& ref_bytes = ref_plane.bytes();
          for (uint32_t j = ref_bytes.a(); j < ref_bytes.b(); ++j) {
            uint32_t r = j % ref_plane.row();
            uint32_t h = j / ref_plane.row();
            if (r < ref_plane.width()) {
              auto diff = ((int16_t)ref_bytes(j) - (int16_t)test_plane(h)(r));
              sse_per_frame += diff * diff;
            }
          }
          num_pixels_per_frame += (ref_plane.width() * ref_plane.height());
        }
        THROW_IF(num_pixels_per_frame == 0, Invalid);
      }
      if (config.verbose) {
        cout << "FRAME " << std::setw(5) << i << " : ";
        double mse_per_frame = sse_per_frame ? (double)sse_per_frame / num_pixels_per_frame : 0;
        cout << calculate_and_get_psnr_string(mse_per_frame) << endl;
      } else {
        cout << "PROCESSING " << std::setw(5) << i << " / " << index1_to_index2.size();
        cout.flush();
        cout << "\r";
      }

      total_sse += sse_per_frame;
      total_num_pixels += num_pixels_per_frame;
    }

    if (config.verbose) {
      cout << "------------------------------" << endl;
      cout << "AVERAGE";
      cout << "    ";
      cout << " : ";
    }
    double mse = total_sse ? total_sse / total_num_pixels : 0;
    cout << calculate_and_get_psnr_string(mse);
    if (!config.verbose) {
      for (auto i = 0; i < 40; ++i) {
        cout << " ";  // clear end of line
      }
    }
    cout << endl;
  } __catch (std::exception& e) {
    cerr << "Error calculating PSNR: " << e.what() << endl;
    return 1;
  }
  return 0;
}