Feat fix megapose model download and add options to MBT tutorial

script/megapose_server/install.py

@@ -107,13 +107,32 @@ def download_models(megapose_env: str, megapose_path: Path, megapose_data_path:
   Download the megapose deep learning models
   '''
   models_path = megapose_data_path / 'megapose-models'
-  conf_path = megapose_path / 'rclone.conf'
-  rclone = str(get_rclone_for_conda_env(megapose_env).absolute())
-  arguments = [rclone, 'copyto', 'inria_data:megapose-models/',
-                str(models_path), '--exclude', '*epoch*',
-                '--config', str(conf_path), '--progress']
-  print(' '.join(arguments))
-  subprocess.run(arguments, check=True)
+  models_path.mkdir(exist_ok=True)
+  try:
+    conf_path = megapose_path / 'rclone.conf'
+    rclone = str(get_rclone_for_conda_env(megapose_env).absolute())
+    arguments = [rclone, 'copyto', 'inria_data:megapose-models/',
+                  str(models_path), '--exclude', '*epoch*',
+                  '--config', str(conf_path), '--progress']
+    print(' '.join(arguments))
+    subprocess.run(arguments, check=True)
+  except:
+    print('Could not download MegaPose data from the original repo, trying to fetch from the ViSP website')
+    from urllib.request import urlretrieve
+
+    base_url = 'https://visp-doc.inria.fr/download/model-zoo/megapose-models/'
+    dirs = ['coarse-rgb-906902141/', 'refiner-rgb-653307694/', 'refiner-rgbd-288182519/']
+    files_in_each_dir = ['checkpoint.pth.tar', 'config.yaml', 'log.txt']
+    for folder_name in dirs:
+      dir_url = base_url + folder_name
+      save_dir = models_path / folder_name
+      save_dir.mkdir(exist_ok=True)
+      for file_name in files_in_each_dir:
+        full_url = dir_url + file_name
+        print(full_url)
+        _, headers = urlretrieve(full_url, str(save_dir / file_name))
+
+
 
 
 def install_server(megapose_env: str):

tutorial/tracking/model-based/generic-rgbd-blender/tutorial-mb-generic-tracker-rgbd-blender.cpp

@@ -86,8 +86,8 @@
 {
   std::cout << "Synopsis" << std::endl
     << "  " << argv[0]
-    << " [--data-path <path>] [--model-path <path>] [--first-frame <index>] [--disable-depth] "
-    << " [--disable-klt] [--step-by-step] [--display-ground-truth] [--help, -h]" << std::endl
+    << " [--data-path <path>] [--model-path <path>] [--first-frame <index>] [--depth-dense-mode <0|1>] "
+    << " [--depth-normals-mode <0|1>] [--me-mode <0|1>] [--klt-mode <0|1>] [--step-by-step] [--display-ground-truth] [--help, -h]" << std::endl
     << std::endl;
   std::cout << "Description" << std::endl
     << "  --data-path <path>   Path to the data generated by Blender get_camera_pose_teabox.py" << std::endl
@@ -100,10 +100,16 @@
     << "  --first-frame <index>   First frame number to process." << std::endl
     << "    Default: " << first_frame << std::endl
     << std::endl
-    << "  --disable-depth  Flag to turn off tracker depth features." << std::endl
+
+    << "  --depth-dense-mode  Whether to use dense depth features (0 = off, 1 = on). default: 1" << std::endl
+    << std::endl
+    << "  --depth-normals-mode  Whether to use normal depth features (0 = off, 1 = on). default: 0" << std::endl
+    << std::endl
+    << "  --me-mode  Whether to use moving edge features (0 = off, 1 = on). default: 1" << std::endl
     << std::endl
-    << "  --disable-klt  Flag to turn off tracker keypoints features." << std::endl
+    << "  --klt-mode  Whether to use KLT features (0 = off, 1 = on). Requires OpenCV. default: 1" << std::endl
     << std::endl
+
     << "  --step-by-step  Flag to enable step by step mode." << std::endl
     << std::endl
     << "  --display-ground-truth  Flag to enable displaying ground truth." << std::endl
@@ -124,8 +130,11 @@
   std::string opt_data_path = "data/teabox";
   std::string opt_model_path = "model/teabox";
   unsigned int opt_first_frame = 1;
-  bool opt_disable_depth = false;
+  int opt_meMode = 1, opt_kltMode = 1, opt_normalsMode = 0, opt_denseMode = 1;
+
+  bool disable_depth = false;
   bool opt_disable_klt = false;
+
   bool opt_display_ground_truth = false;
   bool opt_step_by_step = false;
 
@@ -138,11 +147,37 @@
       opt_model_path = std::string(argv[i + 1]);
       i++;
     }
-    else if (std::string(argv[i]) == "--disable-depth") {
-      opt_disable_depth = true;
+    else if (std::string(argv[i]) == "--depth-dense-mode" && i + 1 < argc) {
+      opt_denseMode = static_cast<unsigned int>(atoi(argv[i + 1]));
+      if (opt_denseMode < 0 || opt_denseMode > 1) {
+        usage(argv, 0, opt_data_path, opt_model_path, opt_first_frame);
+        return EXIT_FAILURE;
+      }
+      i++;
     }
-    else if (std::string(argv[i]) == "--disable-klt") {
-      opt_disable_klt = true;
+    else if (std::string(argv[i]) == "--depth-normals-mode" && i + 1 < argc) {
+      opt_normalsMode = static_cast<unsigned int>(atoi(argv[i + 1]));
+      if (opt_normalsMode < 0 || opt_normalsMode > 1) {
+        usage(argv, 0, opt_data_path, opt_model_path, opt_first_frame);
+        return EXIT_FAILURE;
+      }
+      i++;
+    }
+    else if (std::string(argv[i]) == "--me-mode" && i + 1 < argc) {
+      opt_meMode = static_cast<unsigned int>(atoi(argv[i + 1]));
+      if (opt_meMode < 0 || opt_meMode > 1) {
+        usage(argv, 0, opt_data_path, opt_model_path, opt_first_frame);
+        return EXIT_FAILURE;
+      }
+      i++;
+    }
+    else if (std::string(argv[i]) == "--klt-mode" && i + 1 < argc) {
+      opt_kltMode = static_cast<unsigned int>(atoi(argv[i + 1]));
+      if (opt_kltMode < 0 || opt_kltMode > 1) {
+        usage(argv, 0, opt_data_path, opt_model_path, opt_first_frame);
+        return EXIT_FAILURE;
+      }
+      i++;
     }
     else if (std::string(argv[i]) == "--display-ground-truth") {
       opt_display_ground_truth = true;
@@ -164,6 +199,8 @@
     }
   }
 
+  disable_depth = opt_denseMode == 0 && opt_normalsMode == 0;
+
   std::string video_color_images = vpIoTools::createFilePath(opt_data_path, "color/%04d_L.jpg");
   std::string video_depth_images = vpIoTools::createFilePath(opt_data_path, "depth/Image%04d_R.exr");
   std::string ground_truth = vpIoTools::createFilePath(opt_data_path, "ground-truth/Camera_L_%04d.txt");
@@ -179,10 +216,10 @@
 
   std::cout << "Input data" << std::endl;
   std::cout << "  Color images    : " << video_color_images << std::endl;
-  std::cout << "  Depth images    : " << (opt_disable_depth ? "Disabled" : video_depth_images) << std::endl;
-  std::cout << "  Extrinsics      : " << (opt_disable_depth ? "Disabled" : extrinsic_file) << std::endl;
+  std::cout << "  Depth images    : " << (disable_depth ? "Disabled" : video_depth_images) << std::endl;
+  std::cout << "  Extrinsics      : " << (disable_depth ? "Disabled" : extrinsic_file) << std::endl;
   std::cout << "  Color intrinsics: " << color_intrinsic_file << std::endl;
-  std::cout << "  Depth intrinsics: " << (opt_disable_depth ? "Disabled" : depth_intrinsic_file) << std::endl;
+  std::cout << "  Depth intrinsics: " << (disable_depth ? "Disabled" : depth_intrinsic_file) << std::endl;
   std::cout << "  Ground truth    : " << ground_truth << std::endl;
   std::cout << "Tracker settings" << std::endl;
   std::cout << "  Color config    : " << mbt_config_color << std::endl;
@@ -195,28 +232,52 @@
   }
   else {
     std::cout << "  Init file       : " << mbt_init_file << std::endl;
-    std::cout << "  Features        : moving-edges " << (opt_disable_klt ? "" : "+ keypoints") << (opt_disable_depth ? "" : " + depth") << std::endl;
+    std::cout << "  Features        : moving-edges " << (opt_disable_klt ? "" : "+ keypoints") << (disable_depth ? "" : " + depth") << std::endl;
   }
 
+
   std::vector<int> tracker_types;
-  if (opt_disable_klt) {
-    tracker_types.push_back(vpMbGenericTracker::EDGE_TRACKER);
+  int colorTracker = 0;
+
+
+  if (opt_meMode == 1) {
+    colorTracker |= vpMbGenericTracker::EDGE_TRACKER;
   }
-  else {
+  if (opt_kltMode == 1) {
 #if defined(VISP_HAVE_OPENCV) && defined(HAVE_OPENCV_IMGPROC) && defined(HAVE_OPENCV_VIDEO)
-    tracker_types.push_back(vpMbGenericTracker::EDGE_TRACKER | vpMbGenericTracker::KLT_TRACKER);
+    colorTracker |= vpMbGenericTracker::KLT_TRACKER;
 #else
-    std::cout << "Warning: keypoints cannot be used as features since ViSP is not build with OpenCV 3rd party" << std::endl;
+    std::cerr << "Warning: keypoints cannot be used as features since ViSP is not built with OpenCV 3rd party" << std::endl;
 #endif
   }
-  if (!opt_disable_depth)
-    tracker_types.push_back(vpMbGenericTracker::DEPTH_DENSE_TRACKER);
+
+  if (colorTracker == 0) {
+    std::cerr << "You should use at least one type of color feature. If OpenCV is not installed, KLT features are disabled" << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  tracker_types.push_back(colorTracker);
+
+  if (!disable_depth) {
+    int depthTracker = 0;
+    if (opt_denseMode == 1) {
+      depthTracker |= vpMbGenericTracker::DEPTH_DENSE_TRACKER;
+    }
+    if (opt_normalsMode == 1) {
+      depthTracker |= vpMbGenericTracker::DEPTH_NORMAL_TRACKER;
+    }
+
+    tracker_types.push_back(depthTracker);
+  }
 
   vpMbGenericTracker tracker(tracker_types);
-  if (!opt_disable_depth)
-    tracker.loadConfigFile(mbt_config_color, mbt_config_depth);
-  else
+  if (!disable_depth) {
+    tracker.loadConfigFile(mbt_config_color, mbt_config_depth, true);
+  }
+  else {
     tracker.loadConfigFile(mbt_config_color);
+
+  }
   tracker.loadModel(mbt_cad_model);
   vpCameraParameters cam_color, cam_depth;
 
@@ -231,20 +292,20 @@
     std::cout << "Cannot found intrinsics for camera " << depth_camera_name << std::endl;
   }
 
-  if (!opt_disable_depth)
+  if (!disable_depth)
     tracker.setCameraParameters(cam_color, cam_depth);
   else
     tracker.setCameraParameters(cam_color);
 
   // Reload intrinsics from tracker (useless)
-  if (!opt_disable_depth)
+  if (!disable_depth)
     tracker.getCameraParameters(cam_color, cam_depth);
   else
     tracker.getCameraParameters(cam_color);
   tracker.setDisplayFeatures(true);
   std::cout << "cam_color:\n" << cam_color << std::endl;
 
-  if (!opt_disable_depth)
+  if (!disable_depth)
     std::cout << "cam_depth:\n" << cam_depth << std::endl;
 
   vpImage<uint16_t> I_depth_raw;
@@ -254,7 +315,7 @@
   vpHomogeneousMatrix cMo_ground_truth;
 
   unsigned int frame_cpt = opt_first_frame;
-  read_data(frame_cpt, video_color_images, video_depth_images, opt_disable_depth, ground_truth,
+  read_data(frame_cpt, video_color_images, video_depth_images, disable_depth, ground_truth,
             I, I_depth_raw, depth_width, depth_height, pointcloud, cam_depth, cMo_ground_truth);
   vpImageConvert::createDepthHistogram(I_depth_raw, I_depth);
 
@@ -267,12 +328,12 @@
 #endif
 
   d1.init(I, 0, 0, "Color image");
-  if (!opt_disable_depth) {
+  if (!disable_depth) {
     d2.init(I_depth, static_cast<int>(I.getWidth()), 0, "Depth image");
   }
 
   vpHomogeneousMatrix depth_M_color;
-  if (!opt_disable_depth) {
+  if (!disable_depth) {
     depth_M_color.load(extrinsic_file);
     tracker.setCameraTransformationMatrix("Camera2", depth_M_color);
     std::cout << "depth_M_color:\n" << depth_M_color << std::endl;
@@ -287,7 +348,7 @@
 
   try {
     bool quit = false;
-    while (!quit && read_data(frame_cpt, video_color_images, video_depth_images, opt_disable_depth,
+    while (!quit && read_data(frame_cpt, video_color_images, video_depth_images, disable_depth,
                               ground_truth, I, I_depth_raw, depth_width, depth_height, pointcloud, cam_depth,
                               cMo_ground_truth)) {
       vpImageConvert::createDepthHistogram(I_depth_raw, I_depth);
@@ -298,7 +359,7 @@
         tracker.initFromPose(I, cMo_ground_truth); // I and I_depth must be the same size when using depth features!
       }
       else {
-        if (!opt_disable_depth) {
+        if (!disable_depth) {
           std::map<std::string, const vpImage<unsigned char> *> mapOfImages;
           std::map<std::string, const std::vector<vpColVector> *> mapOfPointClouds;
           std::map<std::string, unsigned int> mapOfPointCloudWidths;
@@ -320,7 +381,7 @@
       if (!opt_display_ground_truth)
         std::cout << "cMo:\n" << cMo << std::endl;
       std::cout << "cMo ground truth:\n" << cMo_ground_truth << std::endl;
-      if (!opt_disable_depth) {
+      if (!disable_depth) {
         tracker.display(I, I_depth, cMo, depth_M_color * cMo, cam_color, cam_depth, vpColor::red, 2);
         vpDisplay::displayFrame(I_depth, depth_M_color * cMo, cam_depth, 0.05, vpColor::none, 2);
       }