Skip to content

Commit

Permalink
Merge pull request #1428 from SamFlt/feat_fix_megapose_mbt_tuto
Browse files Browse the repository at this point in the history 
Feat fix megapose model download and add options to MBT tutorial
  • Loading branch information
@fspindle
fspindle authored Jun 27, 2024
2 parents 503f563 + a447f83 commit 1e7e8e6
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 39 deletions.
33 changes: 26 additions & 7 deletions script/megapose_server/install.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
125 changes: 93 additions & 32 deletions ...al/tracking/model-based/generic-rgbd-blender/tutorial-mb-generic-tracker-rgbd-blender.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,8 +86,8 @@ void usage(const char **argv, int error, const std::string &data_path, const std
{
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
Expand All @@ -100,10 +100,16 @@ void usage(const char **argv, int error, const std::string &data_path, const std
<< " --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
Expand All @@ -124,8 +130,11 @@ int main(int argc, const char **argv)
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;

Expand All @@ -138,11 +147,37 @@ int main(int argc, const char **argv)
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;
Expand All @@ -164,6 +199,8 @@ int main(int argc, const char **argv)
}
}

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");
Expand All @@ -179,10 +216,10 @@ int main(int argc, const char **argv)

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;
Expand All @@ -195,28 +232,52 @@ int main(int argc, const char **argv)
}
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;

Expand All @@ -231,20 +292,20 @@ int main(int argc, const char **argv)
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;
Expand All @@ -254,7 +315,7 @@ int main(int argc, const char **argv)
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);

Expand All @@ -267,12 +328,12 @@ int main(int argc, const char **argv)
#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;
Expand All @@ -287,7 +348,7 @@ int main(int argc, const char **argv)

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);
Expand All @@ -298,7 +359,7 @@ int main(int argc, const char **argv)
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;
Expand All @@ -320,7 +381,7 @@ int main(int argc, const char **argv)
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);
}
Expand Down

0 comments on commit 1e7e8e6

Please sign in to comment.