diff --git a/.gitignore b/.gitignore index a7f162b..5eb7b7d 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,5 @@ data/ # Ignore the nested Catkin workspace's build directories /pose_ws/build/ -/pose_ws/devel/ \ No newline at end of file +/pose_ws/devel/ +/pose_ws/.catkin_tools diff --git a/.gitmodules b/.gitmodules index 4000ed4..63eb5e6 100644 --- a/.gitmodules +++ b/.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 diff --git a/README.md b/README.md index f4317f5..6a7cddc 100644 --- a/README.md +++ b/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. diff --git a/cnos/src/scripts/inference_custom.py b/cnos/src/scripts/inference_custom.py index cb34e58..71fa184 100644 --- a/cnos/src/scripts/inference_custom.py +++ b/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) \ No newline at end of file + 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, + ) diff --git a/object_models b/object_models index 5c64d2d..735a8b1 160000 --- a/object_models +++ b/object_models @@ -1 +1 @@ -Subproject commit 5c64d2dc07fd2bc8101960cba780e36285b2fe75 +Subproject commit 735a8b1cc58e8eb0977bc3832c48addbf9d9bced diff --git a/pose_ws/src/object_detection_msgs b/pose_ws/src/object_detection_msgs deleted file mode 160000 index 809eca2..0000000 --- a/pose_ws/src/object_detection_msgs +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 809eca2ad4f81430e8cc3af8595d72ce5195eaeb diff --git a/pose_ws/src/pose_estimation_msgs b/pose_ws/src/pose_estimation_msgs new file mode 160000 index 0000000..5097fda --- /dev/null +++ b/pose_ws/src/pose_estimation_msgs @@ -0,0 +1 @@ +Subproject commit 5097fda66025a5dff159360c40c86c2e1a54f6cf diff --git a/real/pose_estimation_server.py b/real/pose_estimation_server.py index fc7e3e7..7b1e461 100644 --- a/real/pose_estimation_server.py +++ b/real/pose_estimation_server.py @@ -2,24 +2,28 @@ from __future__ import annotations -import os -from pathlib import Path - import cv2 -from cv_bridge import CvBridge +import geometry_msgs.msg import numpy as np import nvdiffrast.torch as dr import ros_numpy import rospy +import torch import trimesh -from sensor_msgs.msg import Image, CameraInfo -from std_msgs.msg import Header -from object_detection_msgs.srv import EstimatePose, EstimatePoseRequest, EstimatePoseResponse +from pose_estimation_msgs.srv import ( + DetectObjects, + DetectObjectsRequest, + DetectObjectsResponse, + EstimatePose, + EstimatePoseRequest, + EstimatePoseResponse, +) +from transformers import OwlViTForObjectDetection, OwlViTProcessor from FoundationPose.learning.training.predict_pose_refine import PoseRefinePredictor from FoundationPose.learning.training.predict_score import ScorePredictor -from FoundationPose.Utils import draw_posed_3d_box, draw_xyz_axis from real.important_paths import get_obj_filepath +from real.pose_visualization import PoseVisualization from real.run_spot import FoundationPose from real.segment import CNOS from transform_utils.kinematics import Pose3D @@ -34,34 +38,47 @@ def __init__(self, service_name: str = "/detect_object_pose") -> None: :param service_name: Name of the service's ROS topic """ - self._service = rospy.Service(service_name, EstimatePose, self.handle_pose_estimation) + self._pose_service = rospy.Service(service_name, EstimatePose, self.handle_pose_estimation) + self._detect_srv = rospy.Service( + "detect_object_pixel_xy", + DetectObjects, + self.handle_object_detection, + ) + self.obj_to_track: list[str] = [] self._prev_object_name: str | None = None # Name of the most recent pose-estimated object self._prev_pose_estimate: Pose3D | None = None # Most recent pose estimate self.segmenter_map: dict[str, CNOS] = {} # Maps object names to CNOS segmenters - self.pose_estimator_map: dict[str, FoundationPose] = {} # Maps object names to FP models + self.est: FoundationPose | None = None # Maps object names to the transformation matrices that would center their mesh # Reference: https://trimesh.org/trimesh.bounds.html#trimesh.bounds.oriented_bounds self.to_origin_map: dict[str, np.ndarray] = {} self.bbox_map: dict[str, np.ndarray] = {} # Maps object names to their mesh bounding boxes + self.mesh_map: dict[str, trimesh.Trimesh] = {} # not sure what type mesh files are self.fp_device = "cuda:0" # Presumably the device used by FoundationPose - self.bridge = CvBridge() - self.latest_vis_image = None - self.latest_cam_info = None - # Map camera info and publish image for visualization - rospy.Subscriber("/spot/camera/frontright/camera_info", CameraInfo, self.camera_info_callback) - self.cam_info_pub = rospy.Publisher("/pose_visualization/camera_info", CameraInfo, queue_size=1) - self.vis_pub = rospy.Publisher("/pose_visualization/image", Image, queue_size=1) + self.pose_buffer: dict[str, Pose3D] = {} # record the last pose of each object + self.detection_model = OwlViTForObjectDetection.from_pretrained( + "google/owlvit-large-patch14", + ).to(self.fp_device) + self.detection_model.eval() + self.detection_processor = OwlViTProcessor.from_pretrained("google/owlvit-large-patch14") + self.query_map: dict[str, str] = { + "cyan_cube": "small cyan toy cube", + "expo_spray": "blue and white expo spray bottle", + "cabinet_open": "wooden cabinet with silver handle", + } # hardcoded object name mapping in case Owl-ViT doesn't work with object_name + self.detection_score_threshold = 0.15 # Threshold for object detection score - rospy.loginfo("Finished initializing EstimatePoseServer. Now spinning...") + self.cnos_confidence_threshold = 0.05 # Threshold for CNOS - def camera_info_callback(self, msg): - """Callback to store the latest camera info.""" - self.latest_cam_info = msg + image_size = self.detection_model.config.vision_config.image_size + self.pose_visualization = PoseVisualization(image_size, self.detection_score_threshold) + + rospy.loginfo("Finished initializing EstimatePoseServer. Now spinning...") def init_segmenter(self, object_name: str) -> None: """Initialize an object segmenter for the named object, if one does not already exist. @@ -73,24 +90,16 @@ def init_segmenter(self, object_name: str) -> None: self.segmenter_map[object_name] = CNOS(object_name) rospy.loginfo("CNOS model initialized.") - def init_pose_estimator(self, object_name: str) -> None: - """Initialize a pose estimation model for the named object, if one does not already exist. + def init_pose_estimator(self) -> None: + """Initialize a pose estimation model with empty meshes. + + Actual objects will be swapped in by est.reset_object() from a dictionary. :param object_name: Name of the object for which to initialize a pose estimator """ - if object_name not in self.pose_estimator_map: - rospy.loginfo(f"Initializing FoundationPose model for object '{object_name}'...") - - mesh_path = get_obj_filepath(object_name) - mesh = trimesh.load(mesh_path) - - to_origin, extents = trimesh.bounds.oriented_bounds(mesh) - assert to_origin.shape == (4, 4), "Expected a 4x4 transformation matrix!" - assert extents.shape == (3,), "Expected a 3-vector of object extents!" - self.to_origin_map[object_name] = to_origin - - bbox = np.stack([-extents / 2, extents / 2], axis=0).reshape(2, 3) - self.bbox_map[object_name] = bbox + if self.est is None: + rospy.loginfo("Initializing FoundationPose model'...") + mesh_tmp = trimesh.primitives.Box(extents=np.ones(3), transform=np.eye(4)).to_mesh() # TODO: Can we reuse these three next pieces across different objects? rospy.loginfo("Initializing FoundationPose model -- ScorePredictor") @@ -101,93 +110,112 @@ def init_pose_estimator(self, object_name: str) -> None: refiner = PoseRefinePredictor() refiner.device = self.fp_device + # Not sure if we need re-init glctx everytime rospy.loginfo("Initializing FoundationPose model -- RasterizeCudaContext") glctx = dr.RasterizeCudaContext(self.fp_device) - # TODO: What is this debug stuff? Do we care? Is this needed per-object? - debug = 1 - debug_dir = Path(__file__).parent / "debug" # Maybe append object name? - os.system( - f"rm -rf {debug_dir}/* && mkdir -p {debug_dir}/track_vis {debug_dir}/ob_in_cam", - ) - rospy.loginfo("Initializing FoundationPose model -- FoundationPose") estimator = FoundationPose( - model_pts=mesh.vertices, - model_normals=mesh.vertex_normals, - mesh=mesh, + model_pts=mesh_tmp.vertices.copy(), + model_normals=mesh_tmp.vertex_normals.copy(), + mesh=mesh_tmp, scorer=scorer, refiner=refiner, - debug_dir=debug_dir, - debug=debug, glctx=glctx, ) # TODO: Maybe adjust below section to also apply when object already has an estimator? - # estimator.to_device("cuda:0") estimator.scorer.model.to(self.fp_device) estimator.refiner.model.to(self.fp_device) estimator.device = self.fp_device - rospy.loginfo(f"FoundationPose model initialized for object '{object_name}'.") + rospy.loginfo("FoundationPose model initialized.") - self.pose_estimator_map[object_name] = estimator - else: - rospy.loginfo(f"FoundationPose model for object '{object_name}' already initialized.") + self.est = estimator - # TODO: Anything needed to re-initialize if another object was pose-estimated since? - - def visualize_pose( - self, - object_name: str, - pose: np.ndarray, - rgb_image: np.ndarray, - camera_k_matrix: np.ndarray, - ) -> None: - """Visualize the given pose estimate for the named object. + def get_mesh(self, object_name): + """Retrive mesh regarding an obj if stored before, otherwise read, calculate, + and store mesh, to_origin, and bbox. :param object_name: Name of the object corresponding to the pose estimate - :param pose: Estimated pose of the object as a homogeneous transformation matrix - :param rgb_image: RGB image from which the pose was estimated - :param camera_k_matrix: Intrinsic camera matrix (K) """ - to_origin = self.to_origin_map[object_name] - bbox = self.bbox_map[object_name] - center_pose = pose @ np.linalg.inv(to_origin) - - vis = draw_posed_3d_box(K=camera_k_matrix, img=rgb_image, ob_in_cam=center_pose, bbox=bbox) - vis = draw_xyz_axis( - color=rgb_image, - ob_in_cam=center_pose, - scale=0.1, - K=camera_k_matrix, - thickness=2, - transparency=0.7, - is_input_rgb=True, - ) + if object_name not in self.mesh_map: + mesh_path = get_obj_filepath(object_name) + mesh = trimesh.load(mesh_path) + self.mesh_map[object_name] = mesh + to_origin, extents = trimesh.bounds.oriented_bounds(mesh) + assert to_origin.shape == (4, 4), "Expected a 4x4 transformation matrix!" + assert extents.shape == (3,), "Expected a 3-vector of object extents!" + self.to_origin_map[object_name] = to_origin + + bbox = np.stack([-extents / 2, extents / 2], axis=0).reshape(2, 3) + self.bbox_map[object_name] = bbox + else: + mesh = self.mesh_map[object_name] + to_origin = self.to_origin_map[object_name] + bbox = self.to_origin_map[object_name] + + return mesh + + def handle_object_detection(self, request: DetectObjectsRequest) -> DetectObjectsResponse: + """Handle a service request to run object detection.""" + image = ros_numpy.numpify(request.image) + # TODO: Maybe need PILImage conversion too + + response = DetectObjectsResponse() + + for text_query in request.text_queries: + # Process image and text inputs + inputs = self.detection_processor( + text=[text_query], + images=image, + return_tensors="pt", + ).to(self.fp_device) + + # Get predictions + with torch.no_grad(): + outputs = self.detection_model(**inputs) + + logits = torch.max(outputs["logits"][0], dim=-1) + scores = torch.sigmoid(logits.values).cpu().detach().numpy() + + # labels = logits.indices.cpu().detach().numpy() # no need to use labels + boxes = outputs["pred_boxes"][0].cpu().detach().numpy() + score, box = max( + [(score, box) for score, box in zip(scores, boxes)], + key=lambda x: x[0], + ) + cx, cy, bw, bh = box + + # owl-vit resize the model, so we need to resize it back with the box + original_y, original_x, _ = image.shape + cx_new = cx * original_x + cy_new = cy * original_y + + pixel = (int(cx_new), int(cy_new)) + + print(f"Picking point: {pixel}", pixel) + print("size of the image", original_x, original_y) + + cv2.circle(image, pixel, 30, (255, 0, 0), 5) + cv2.imshow("Stop if bad (press any key to exit)", image) + cv2.waitKey() + cv2.destroyAllWindows() + + pixel_msg = geometry_msgs.msg.Point(pixel[0], pixel[1], 0) + + response.pixels.append(pixel_msg) + + rospy.loginfo("Exiting handle_object_detection()...") + return response - self.latest_vis_image = vis - - # Publish the latest visualization and camera info with the original frame_id - ros_image = self.bridge.cv2_to_imgmsg(self.latest_vis_image, encoding="passthrough") - ros_image.header = Header() - ros_image.header.stamp = rospy.Time.now() - ros_image.header.frame_id = "frontright_fisheye" - self.vis_pub.publish(ros_image) - - if self.latest_cam_info: - cam_info_msg = self.latest_cam_info - cam_info_msg.header.stamp = rospy.Time.now() - cam_info_msg.header.frame_id = "frontright_fisheye" - self.cam_info_pub.publish(cam_info_msg) - def handle_pose_estimation(self, request: EstimatePoseRequest) -> EstimatePoseResponse: """Handle a service request to estimate an object's pose. :param request: Request message for the EstimatePose service """ - object_name = request.query + object_name = request.object_name rospy.loginfo(f"Received request for pose estimation of object '{object_name}'...") rgb_image = ros_numpy.numpify(request.rgb) @@ -197,18 +225,50 @@ def handle_pose_estimation(self, request: EstimatePoseRequest) -> EstimatePoseRe camera_k_matrix = np.array(request.info.K).reshape(3, 3) # If necessary, initialize CNOS and FoundationPose for the requested object + if object_name not in self.segmenter_map: + self.init_segmenter(object_name) + self.init_pose_estimator() + assert self.est is not None, "Expected FoundationPose estimator to exist by now!" + + # call Owl-ViT and check the score. Use the object's predefined text prompt if it exists + text_query = self.query_map.get(object_name, [object_name]) + rospy.loginfo(f"Running object detection with query: '{text_query}'") + + detection_query = self.detection_processor( + text=text_query, + images=rgb_image, + return_tensors="pt", + ).to(self.fp_device) + + with torch.no_grad(): + detection_outputs = self.detection_model(**detection_query) + + # Get prediction logits + logits = torch.max(detection_outputs["logits"][0], dim=-1) + detection_scores = torch.sigmoid(logits.values).cpu().detach().numpy() + + if max(detection_scores) > self.detection_score_threshold: + print(f"Object '{object_name}' had maximum detection score: {max(detection_scores)}.") + + mesh = self.get_mesh(object_name) + # NOTE: not sure what symmetry_tf is in the original code. Omit it for now + # https://github.com/NVlabs/FoundationPose/blob/main/run_ycb_video.py#L118C96-L118C108 + self.est.reset_object( + model_pts=mesh.vertices.copy(), + model_normals=mesh.vertex_normals.copy(), + mesh=mesh, + ) - self.init_segmenter(object_name) - self.init_pose_estimator(object_name) - - # Check whether we can reuse the most recent pose estimate (if not, run FastSAM) - if self._prev_pose_estimate is None or object_name != self._prev_object_name: - segmenter = self.segmenter_map[object_name] - masks = segmenter.run_inference(rgb_image, num_max_dets=1, conf_threshold=0.0) + # NOTE: registration only, no tracking for now! + masks = self.segmenter_map[object_name].run_inference( + rgb_image, + num_max_dets=1, + conf_threshold=self.cnos_confidence_threshold, + ) mask = masks[0] - # TODO: Unsure what the following comment meant: "call this after resetting obj model" - pose_estimate_matrix = self.pose_estimator_map[object_name].register( + # Reinitialize the pose estimator using register() after switching to a new object + pose_estimate_matrix = self.est.register( K=camera_k_matrix, rgb=rgb_image, depth=depth_image, @@ -216,32 +276,57 @@ def handle_pose_estimation(self, request: EstimatePoseRequest) -> EstimatePoseRe iteration=5, ) - else: # TODO: Could/should we re-use non-consecutive pose estimates for the same object? - pose_estimate_matrix = self.pose_estimator_map[object_name].track_one( - rgb=rgb_image, - depth=depth_image, - K=camera_k_matrix, - iteration=2, + if hasattr(self.est, "poses"): + print(f"FoundationPose has {len(self.est.poses)} poses: {self.est.poses}") + + confidence = 0.0 + if hasattr(self.est, "scores"): + print(f"FoundationPose has {len(self.est.scores)} scores: {self.est.scores}") + confidence = self.est.scores[0] + rospy.loginfo(f"Max. confidence was: {confidence:.2f}") + + # Estimated pose gives the frame of the object (o) w.r.t. the camera frame (c) + estimated_pose_c_o = Pose3D.from_homogeneous_matrix(pose_estimate_matrix) + rospy.loginfo(f"Estimated pose: {estimated_pose_c_o}...") + + to_origin = self.to_origin_map[object_name] + center_pose = pose_estimate_matrix @ np.linalg.inv(to_origin) + + # Visualize the estimated pose for the requested object + self.pose_visualization.visualize_pose( + object_name=object_name, + center_pose=center_pose, + bbox=self.bbox_map[object_name], + camera_info=request.info, + rgb_image=rgb_image, + camera_k_matrix=camera_k_matrix, + detection_outputs=detection_outputs, ) - # Estimated pose gives the frame of the object (o) w.r.t. the camera frame (c) - estimated_pose_c_o = Pose3D.from_homogeneous_matrix(pose_estimate_matrix) - rospy.loginfo(f"Estimated pose: {estimated_pose_c_o}...") + # Finally, return a response for the ROS service + response = EstimatePoseResponse() + + # TODO: check self.scores to determine the condition to stop publishing estimated pose + response.object_found = True + response.pose = pose_to_stamped_msg(estimated_pose_c_o) + response.pose.header.stamp = rospy.Time.now() + response.confidence = confidence - # Visualize the estimated pose for the requested object - self.visualize_pose(object_name, pose_estimate_matrix, rgb_image, camera_k_matrix) + self.pose_buffer[object_name] = estimated_pose_c_o - # Update stored most-recent pose estimate and corresponding object name - self._prev_object_name = object_name - self._prev_pose_estimate = estimated_pose_c_o + rospy.loginfo(f"Responding to the EstimatePose client for object {object_name}...") - # Finally, return a response for the ROS service - response = EstimatePoseResponse() - response.object_found = True # TODO: How does/can FoundationPose indicate no object found? - response.pose = pose_to_stamped_msg(estimated_pose_c_o) - response.pose.header.stamp = rospy.Time.now() + else: + response = EstimatePoseResponse() + + response.object_found = False + if object_name in self.pose_buffer: + response.pose = pose_to_stamped_msg(self.pose_buffer[object_name]) + response.pose.header.stamp = rospy.Time.now() - rospy.loginfo("Sending response back to the EstimatePose client...") + rospy.loginfo( + f"Object not found! Retrieving last estimated pose for object {object_name}...", + ) return response @@ -249,7 +334,7 @@ def handle_pose_estimation(self, request: EstimatePoseRequest) -> EstimatePoseRe def main() -> None: """Launch a server for the DetectObjectPose ROS service.""" rospy.init_node("pose_estimation_server") - service = EstimatePoseServer() + _ = EstimatePoseServer() rospy.spin() diff --git a/real/pose_visualization.py b/real/pose_visualization.py new file mode 100644 index 0000000..9e098fe --- /dev/null +++ b/real/pose_visualization.py @@ -0,0 +1,135 @@ +"""Define a utility class to visualize estimated poses.""" + +from __future__ import annotations + +import cv2 +import numpy as np +import rospy +import sensor_msgs.msg +import torch +from cv_bridge import CvBridge +from PIL import Image +from sensor_msgs.msg import Image as ImageMsg +from transformers.image_utils import ImageFeatureExtractionMixin + +from FoundationPose.Utils import draw_posed_3d_box, draw_xyz_axis + + +class PoseVisualization: + """A wrapper for ROS publishers used to visualize estimated poses.""" + + def __init__(self, image_size, score_threshold: float) -> None: + """Initialize the ROS publishers needed to visualize poses.""" + self.image_size = image_size + self.score_threshold = score_threshold + + self.latest_image: dict[str, np.ndarray] = {} # Map camera names to their latest images + + # Publish each camera's latest image to /pose_visualization/[camera_name]/image + self.per_camera_pubs: dict[str, rospy.Publisher] = {} + + self.camera_info_pub = rospy.Publisher( + "/pose_visualization/camera_info", + sensor_msgs.msg.CameraInfo, + queue_size=1, + ) + + self.mixin = ImageFeatureExtractionMixin() # visualization tool + self.bridge = CvBridge() + + def get_publisher(self, camera_name: str) -> rospy.Publisher: + """Retrieve the publisher for the named camera, creating it if necessary. + + :param camera_name: Name of the camera that took the image to be published + :return: Image publisher with topic named after the camera + """ + if camera_name not in self.per_camera_pubs: + topic_name = f"/pose_visualization/{camera_name}/image" + self.per_camera_pubs[camera_name] = rospy.Publisher(topic_name, ImageMsg, queue_size=1) + + return self.per_camera_pubs[camera_name] + + def draw_predictions(self, image, text_queries, scores, boxes, labels) -> np.ndarray: + original_size = Image.fromarray(image, "RGB").size + rospy.loginfo(original_size) + + image = self.mixin.resize(image, self.image_size) + w, h = image.size + image_copy = np.asarray(image).copy() + + for score, box, label in zip(scores, boxes, labels): + if score < self.score_threshold and score < max(scores): + continue + cx, cy, bw, bh = box + x1, y1 = int((cx - bw / 2) * w), int((cy - bh / 2) * h) + x2, y2 = int((cx + bw / 2) * w), int((cy + bh / 2) * h) + + # Draw bounding box + cv2.rectangle(image_copy, (x1, y1), (x2, y2), (255, 0, 0), 2) + + # Add label text + label_text = f"{text_queries[label]}: {score:.2f}" + cv2.putText( + image_copy, + label_text, + (x1, y1 + 15), + cv2.FONT_HERSHEY_SIMPLEX, + 0.5, + (0, 255, 0), + 1, + cv2.LINE_AA, + ) + # resize it back + image_restored = self.mixin.resize(image_copy, original_size) + return np.asarray(image_restored) + + def visualize_pose( + self, + object_name: str, + center_pose: np.ndarray, + bbox: np.ndarray, + camera_info: sensor_msgs.msg.CameraInfo, + rgb_image: np.ndarray, + camera_k_matrix: np.ndarray, + detection_outputs, + ) -> None: + """Visualize the given pose estimate for the named object. + + :param object_name: Name of the object corresponding to the pose estimate + :param center_pose: Estimated pose of the object as a transform matrix (TODO: what frame?) + :param bbox: TODO + :param camera_info: ROS message characterizing the camera that took this image + :param rgb_image: RGB image from which the pose was estimated + :param camera_k_matrix: Intrinsic camera matrix (K) + """ + vis = draw_posed_3d_box(K=camera_k_matrix, img=rgb_image, ob_in_cam=center_pose, bbox=bbox) + vis = draw_xyz_axis( + color=rgb_image, + ob_in_cam=center_pose, + scale=0.1, + K=camera_k_matrix, + thickness=2, + transparency=0.7, + is_input_rgb=True, + ) + + # Draw 2D bounding box for object detection + # Get prediction logits, labels, and boundary boxes + logits = torch.max(detection_outputs["logits"][0], dim=-1) + scores = torch.sigmoid(logits.values).cpu().detach().numpy() + labels = logits.indices.cpu().detach().numpy() + boxes = detection_outputs["pred_boxes"][0].cpu().detach().numpy() + vis = self.draw_predictions(vis, [object_name], scores, boxes, labels) + + camera = camera_info.header.frame_id + image_pub = self.get_publisher(camera) + self.latest_image[camera] = self.draw_predictions(vis, [object_name], scores, boxes, labels) + + # Publish the latest visualization and camera info with the original frame_id + ros_image = self.bridge.cv2_to_imgmsg(self.latest_image[camera], encoding="passthrough") + ros_image.header.stamp = rospy.Time.now() + ros_image.header.frame_id = camera_info.header.frame_id + image_pub.publish(ros_image) + + # Don't overwrite the camera info timestamp, which gives when the original image was taken + self.camera_info_pub.publish(camera_info) diff --git a/run_service.sh b/run_service.sh index f7fc300..2398af6 100644 --- a/run_service.sh +++ b/run_service.sh @@ -2,6 +2,7 @@ cd /docker/pose pip install -e . +pip install transformers source pose_ws/devel/setup.bash python real/pose_estimation_server.py diff --git a/transform_utils b/transform_utils index 75de695..7c626c2 160000 --- a/transform_utils +++ b/transform_utils @@ -1 +1 @@ -Subproject commit 75de6951b754341630aa5b1abad50b3cb8baeb66 +Subproject commit 7c626c2bedcb849181884c46f00fcb51710e2441