Add Support for Multi-Object Pose Estimation

.gitignore

@@ -25,4 +25,5 @@ data/
 
 # Ignore the nested Catkin workspace's build directories
 /pose_ws/build/
-/pose_ws/devel/
+/pose_ws/devel/
+/pose_ws/.catkin_tools

.gitmodules

@@ -1,9 +1,9 @@
-[submodule "pose_ws/src/object_detection_msgs"]
-	path = pose_ws/src/object_detection_msgs
-	url = git@github.com:Benned-H/object_detection_msgs.git
 [submodule "transform_utils"]
 	path = transform_utils
 	url = git@github.com:Benned-H/transform_utils.git
 [submodule "object_models"]
 	path = object_models
 	url = git@github.com:h2r/object_models.git
+[submodule "pose_ws/src/pose_estimation_msgs"]
+	path = pose_ws/src/pose_estimation_msgs
+	url = git@github.com:h2r/pose_estimation_msgs.git

README.md

@@ -1,6 +1,6 @@
 # Pose Estimation ROS Workspace
 
-This repository is structured as a Catkin/ROS workspace containing a single Catkin package (`object_detection_msgs`) used to build ROS service messages. Alongside the typical Catkin folder structure, we have changed the repository structure as little as possible from the original setup.
+This repository is structured as a Catkin/ROS workspace containing a single Catkin package (`pose_estimation_msgs`) used to build ROS service messages. Alongside the typical Catkin folder structure, we have changed the repository structure as little as possible from the original setup.
 
 **NOTE**: The following commands have not been verified to work until this comment is removed.
 

cnos/src/scripts/inference_custom.py

@@ -1,41 +1,45 @@
-import os, sys
-import numpy as np
+import logging
+import os
+import os.path as osp
 import shutil
-from tqdm import tqdm
+import sys
 import time
+
+import numpy as np
 import torch
+from hydra import compose, initialize
 from PIL import Image
-import logging
-import os, sys
-import os.path as osp
-from hydra import initialize, compose
+from tqdm import tqdm
+
 # set level logging
 logging.basicConfig(level=logging.INFO)
-import logging
-import numpy as np
-from hydra.utils import instantiate
 import argparse
 import glob
-from src.utils.bbox_utils import CropResizePad
-from omegaconf import DictConfig, OmegaConf
-from torchvision.utils import save_image
-import torchvision.transforms as T
-from src.model.utils import Detections, convert_npz_to_json
-from src.model.loss import Similarity
-from src.utils.inout import save_json_bop23
+import logging
+
 import cv2
 import distinctipy
+import torchvision.transforms as T
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from segment_anything.utils.amg import rle_to_mask
 from skimage.feature import canny
 from skimage.morphology import binary_dilation
-from segment_anything.utils.amg import rle_to_mask
+from src.model.loss import Similarity
+from src.model.utils import Detections, convert_npz_to_json
+from src.utils.bbox_utils import CropResizePad
+from src.utils.inout import save_json_bop23
+from torchvision.utils import save_image
+
 inv_rgb_transform = T.Compose(
-        [
-            T.Normalize(
-                mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
-                std=[1 / 0.229, 1 / 0.224, 1 / 0.225],
-            ),
-        ]
-    )
+    [
+        T.Normalize(
+            mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
+            std=[1 / 0.229, 1 / 0.224, 1 / 0.225],
+        ),
+    ],
+)
+
 
 def visualize(rgb, detections, save_path="./tmp/tmp.png"):
     img = rgb.copy()
@@ -52,28 +56,29 @@ def visualize(rgb, detections, save_path="./tmp/tmp.png"):
         obj_id = det["category_id"]
         temp_id = obj_id - 1
 
-        r = int(255*colors[temp_id][0])
-        g = int(255*colors[temp_id][1])
-        b = int(255*colors[temp_id][2])
-        img[mask, 0] = alpha*r + (1 - alpha)*img[mask, 0]
-        img[mask, 1] = alpha*g + (1 - alpha)*img[mask, 1]
-        img[mask, 2] = alpha*b + (1 - alpha)*img[mask, 2]   
+        r = int(255 * colors[temp_id][0])
+        g = int(255 * colors[temp_id][1])
+        b = int(255 * colors[temp_id][2])
+        img[mask, 0] = alpha * r + (1 - alpha) * img[mask, 0]
+        img[mask, 1] = alpha * g + (1 - alpha) * img[mask, 1]
+        img[mask, 2] = alpha * b + (1 - alpha) * img[mask, 2]
         img[edge, :] = 255
-    
+
     img = Image.fromarray(np.uint8(img))
     img.save(save_path)
     prediction = Image.open(save_path)
-    
+
     # concat side by side in PIL
     img = np.array(img)
-    concat = Image.new('RGB', (img.shape[1] + prediction.size[0], img.shape[0]))
+    concat = Image.new("RGB", (img.shape[1] + prediction.size[0], img.shape[0]))
     concat.paste(rgb, (0, 0))
     concat.paste(prediction, (img.shape[1], 0))
     return concat
-
+
+
 def run_inference(template_dir, rgb_path, num_max_dets, conf_threshold, stability_score_thresh):
     with initialize(version_base=None, config_path="../../configs"):
-        cfg = compose(config_name='run_inference.yaml')
+        cfg = compose(config_name="run_inference.yaml")
     cfg_segmentor = cfg.model.segmentor_model
     if "fast_sam" in cfg_segmentor._target_:
         logging.info("Using FastSAM, ignore stability_score_thresh!")
@@ -82,20 +87,17 @@ def run_inference(template_dir, rgb_path, num_max_dets, conf_threshold, stabilit
     metric = Similarity()
     logging.info("Initializing model")
     model = instantiate(cfg.model)
-    
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.descriptor_model.model = model.descriptor_model.model.to(device)
     model.descriptor_model.model.device = device
     # if there is predictor in the model, move it to device
     if hasattr(model.segmentor_model, "predictor"):
-        model.segmentor_model.predictor.model = (
-            model.segmentor_model.predictor.model.to(device)
-        )
+        model.segmentor_model.predictor.model = model.segmentor_model.predictor.model.to(device)
     else:
         model.segmentor_model.model.setup_model(device=device, verbose=True)
     logging.info(f"Moving models to {device} done!")
-
-
+
     logging.info("Initializing template")
     template_paths = glob.glob(f"{template_dir}/*.png")
     boxes, templates = [], []
@@ -105,61 +107,88 @@ def run_inference(template_dir, rgb_path, num_max_dets, conf_threshold, stabilit
 
         image = torch.from_numpy(np.array(image.convert("RGB")) / 255).float()
         templates.append(image)
-        
+
     templates = torch.stack(templates).permute(0, 3, 1, 2)
     boxes = torch.tensor(np.array(boxes))
-    
+
     processing_config = OmegaConf.create(
         {
             "image_size": 224,
-        }
+        },
     )
     proposal_processor = CropResizePad(processing_config.image_size)
     templates = proposal_processor(images=templates, boxes=boxes).cuda()
     save_image(templates, f"{template_dir}/cnos_results/templates.png", nrow=7)
     ref_feats = model.descriptor_model.compute_features(
-                    templates, token_name="x_norm_clstoken"
-                )
+        templates,
+        token_name="x_norm_clstoken",
+    )
     logging.info(f"Ref feats: {ref_feats.shape}")
-    
+
     # run inference
     rgb = Image.open(rgb_path).convert("RGB")
     detections = model.segmentor_model.generate_masks(np.array(rgb))
     detections = Detections(detections)
     decriptors = model.descriptor_model.forward(np.array(rgb), detections)
-    
+
     # get scores per proposal
     scores = metric(decriptors[:, None, :], ref_feats[None, :, :])
     score_per_detection = torch.topk(scores, k=5, dim=-1)[0]
     score_per_detection = torch.mean(
-        score_per_detection, dim=-1
+        score_per_detection,
+        dim=-1,
     )
-    
+
     # get top-k detections
     scores, index = torch.topk(score_per_detection, k=num_max_dets, dim=-1)
     detections.filter(index)
-    
+
     # keep only detections with score > conf_threshold
-    detections.filter(scores>conf_threshold)
+    detections.filter(scores > conf_threshold)
     detections.add_attribute("scores", scores)
     detections.add_attribute("object_ids", torch.zeros_like(scores))
-        
+
     detections.to_numpy()
     save_path = f"{template_dir}/cnos_results/detection"
     detections.save_to_file(0, 0, 0, save_path, "custom", return_results=False)
-    detections = convert_npz_to_json(idx=0, list_npz_paths=[save_path+".npz"])
-    save_json_bop23(save_path+".json", detections)
+    detections = convert_npz_to_json(idx=0, list_npz_paths=[save_path + ".npz"])
+    save_json_bop23(save_path + ".json", detections)
     vis_img = visualize(rgb, detections)
     vis_img.save(f"{template_dir}/cnos_results/vis.png")
-
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--template_dir", nargs="?", help="Path to root directory of the template")
     parser.add_argument("--rgb_path", nargs="?", help="Path to RGB image")
-    parser.add_argument("--num_max_dets", nargs="?", default=1, type=int, help="Number of max detections")
-    parser.add_argument("--confg_threshold", nargs="?", default=0.5, type=float, help="Confidence threshold")
-    parser.add_argument("--stability_score_thresh", nargs="?", default=0.97, type=float, help="stability_score_thresh of SAM")
+    parser.add_argument(
+        "--num_max_dets",
+        nargs="?",
+        default=1,
+        type=int,
+        help="Number of max detections",
+    )
+    parser.add_argument(
+        "--confg_threshold",
+        nargs="?",
+        default=0.5,
+        type=float,
+        help="Confidence threshold",
+    )
+    parser.add_argument(
+        "--stability_score_thresh",
+        nargs="?",
+        default=0.97,
+        type=float,
+        help="stability_score_thresh of SAM",
+    )
     args = parser.parse_args()
 
     os.makedirs(f"{args.template_dir}/cnos_results", exist_ok=True)
-    run_inference(args.template_dir, args.rgb_path, num_max_dets=args.num_max_dets, conf_threshold=args.confg_threshold, stability_score_thresh=args.stability_score_thresh)
+    run_inference(
+        args.template_dir,
+        args.rgb_path,
+        num_max_dets=args.num_max_dets,
+        conf_threshold=args.confg_threshold,
+        stability_score_thresh=args.stability_score_thresh,
+    )