vpMegaPose.h

/*
 * ViSP, open source Visual Servoing Platform software.
 * Copyright (C) 2005 - 2025 by Inria. All rights reserved.
 *
 * This software is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * See the file LICENSE.txt at the root directory of this source
 * distribution for additional information about the GNU GPL.
 *
 * For using ViSP with software that can not be combined with the GNU
 * GPL, please contact Inria about acquiring a ViSP Professional
 * Edition License.
 *
 * See https://visp.inria.fr for more information.
 *
 * This software was developed at:
 * Inria Rennes - Bretagne Atlantique
 * Campus Universitaire de Beaulieu
 * 35042 Rennes Cedex
 * France
 *
 * If you have questions regarding the use of this file, please contact
 * Inria at visp@inria.fr
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * Description:
 * MegaPose wrapper.
 */
#ifndef VP_MEGAPOSE_H
#define VP_MEGAPOSE_H
 
#include <visp3/core/vpConfig.h>
#if defined(VISP_HAVE_NLOHMANN_JSON) && defined(VISP_HAVE_THREADS)
 
#include <vector>
#include <string>
#include <unordered_map>
#include <mutex>
 
#include <visp3/core/vpHomogeneousMatrix.h>
#include <visp3/core/vpCameraParameters.h>
#include <visp3/core/vpImage.h>
#include <visp3/core/vpRGBa.h>
#include <visp3/core/vpRect.h>
 
#include VISP_NLOHMANN_JSON(json.hpp)
 
BEGIN_VISP_NAMESPACE
class vpMegaPoseEstimate
{
public:
  vpMegaPoseEstimate() : score(0.f) { }
  vpHomogeneousMatrix cTo;
  float score;
  vpRect boundingBox;
};
 
// Don't use the default ViSP JSON conversion to avoid potential regressions (that wouldn't be detected)
// and very specific parsing on the server side
inline void from_megapose_json(const nlohmann::json &j, vpHomogeneousMatrix &T)
{
  std::vector<double> values = j;
  assert(values.size() == 16);
  std::copy(values.begin(), values.end(), T.data);
}
 
inline void to_megapose_json(nlohmann::json &j, const vpHomogeneousMatrix &T)
{
  std::vector<double> values;
  values.reserve(16);
  for (unsigned i = 0; i < 16; ++i) {
    values.push_back(T.data[i]);
  }
  j = values;
}
 
inline void to_megapose_json(nlohmann::json &j, const vpRect &d)
{
  std::vector<double> values = {
    d.getLeft(), d.getTop(), d.getRight(), d.getBottom()
  };
  j = values;
}
 
inline void from_megapose_json(const nlohmann::json &j, vpRect &d)
{
  std::vector<double> values = j.get<std::vector<double>>();
  assert((values.size() == 4));
  d.setLeft(values[0]);
  d.setTop(values[1]);
  d.setRight(values[2]);
  d.setBottom(values[3]);
}
 
inline void from_json(const nlohmann::json &j, vpMegaPoseEstimate &m)
{
  m.score = j["score"];
  from_megapose_json(j.at("cTo"), m.cTo);
  from_megapose_json(j.at("boundingBox"), m.boundingBox);
}

class VISP_EXPORT vpMegaPose
{
public:
  enum ServerMessage
  {
    UNKNOWN = 0,
    ERR = 1, 
    OK = 2, 
    GET_POSE = 3, 
    RET_POSE = 4, 
    GET_VIZ = 5, 
    RET_VIZ = 6, 
    SET_INTR = 7, 
    GET_SCORE = 8, 
    RET_SCORE = 9, 
    SET_SO3_GRID_SIZE = 10, 
    GET_LIST_OBJECTS = 11,
    RET_LIST_OBJECTS = 12,
    EXIT = 13
  };

  vpMegaPose(const std::string &host, int port, const vpCameraParameters &cam, unsigned height, unsigned width);

  std::vector<vpMegaPoseEstimate> estimatePoses(const vpImage<vpRGBa> &image, const std::vector<std::string> &objectNames,
                                                const vpImage<uint16_t> *const depth = nullptr, const double depthToM = 0.f,
                                                const std::vector<vpRect> *const detections = nullptr,
                                                const std::vector<vpHomogeneousMatrix> *const initial_cTos = nullptr,
                                                int refinerIterations = -1);
  std::vector<double> scorePoses(const vpImage<vpRGBa> &image, const std::vector<std::string> &objectNames,
                                  const std::vector<vpHomogeneousMatrix> &cTos);

  void setIntrinsics(const vpCameraParameters &cam, unsigned height, unsigned width);
 
  vpImage<vpRGBa> viewObjects(const std::vector<std::string> &objectNames,
                              const std::vector<vpHomogeneousMatrix> &poses, const std::string &viewType);

  void setCoarseNumSamples(const unsigned num);

  std::vector<std::string> getObjectNames();
 
  ~vpMegaPose();
 
private:
  // Server connection data
  int m_serverSocket;
  int m_fd;
 
  std::mutex m_mutex; // Since client-server communications are synchronous, avoid multiple parallel communications
 
  void makeMessage(const vpMegaPose::ServerMessage messageType, std::vector<uint8_t> &data) const;
  std::pair<vpMegaPose::ServerMessage, std::vector<uint8_t>> readMessage() const;
 
  const static std::unordered_map<vpMegaPose::ServerMessage, std::string> m_codeMap;
  static std::string messageToString(const vpMegaPose::ServerMessage messageType);
  static vpMegaPose::ServerMessage stringToMessage(const std::string &s);
};
 
END_VISP_NAMESPACE
#endif // VISP_HAVE_NLOHMANN_JSON
#endif