| import argparse | |
| import os | |
| import warnings | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.cuda.amp as amp | |
| import torch.nn.functional as F | |
| from loguru import logger | |
| import utils.config as config | |
| from model import build_net | |
| from utils.grasp_eval import detect_grasps | |
| from utils.misc import setup_logger | |
| from infer_custom_image_nodepth import ( | |
| default_save_path, | |
| inverse_affine, | |
| preprocess_image, | |
| save_prediction_overlay, | |
| ) | |
| warnings.filterwarnings("ignore") | |
| cv2.setNumThreads(0) | |
| ALLOWED_MISSING_PREFIXES = ( | |
| "llava_model.", | |
| "module.llava_model.", | |
| ) | |
| def get_parser(): | |
| parser = argparse.ArgumentParser( | |
| description="Run non-depth MaPLe custom inference on one RGB image and sentence" | |
| ) | |
| parser.add_argument("--config", required=True, type=str, help="config file") | |
| parser.add_argument("--image", required=True, type=str, help="input RGB image path") | |
| parser.add_argument("--sentence", required=True, type=str, help="referring expression") | |
| parser.add_argument("--save", default="", type=str, help="output visualization path") | |
| parser.add_argument("--num-grasps", default=5, type=int, help="number of grasp boxes to draw") | |
| parser.add_argument("--mask-threshold", default=0.35, type=float, help="mask threshold") | |
| parser.add_argument( | |
| "--gate-grasps-by-mask", | |
| action="store_true", | |
| help="also gate final grasp peak detection by the predicted instance mask", | |
| ) | |
| parser.add_argument( | |
| "--opts", | |
| default=None, | |
| nargs=argparse.REMAINDER, | |
| help="override settings in the config.", | |
| ) | |
| cli_args = parser.parse_args() | |
| cfg = config.load_cfg_from_cfg_file(cli_args.config) | |
| if cli_args.opts is not None: | |
| cfg = config.merge_cfg_from_list(cfg, cli_args.opts) | |
| cfg.custom_image = cli_args.image | |
| cfg.custom_sentence = cli_args.sentence | |
| cfg.custom_save = cli_args.save | |
| cfg.custom_num_grasps = cli_args.num_grasps | |
| cfg.custom_mask_threshold = cli_args.mask_threshold | |
| cfg.custom_gate_grasps_by_mask = cli_args.gate_grasps_by_mask | |
| return cfg | |
| def _normalize_state_dict_for_model(model, state_dict): | |
| model_keys = model.state_dict().keys() | |
| ckpt_keys = state_dict.keys() | |
| model_has_module = len(model_keys) > 0 and next(iter(model_keys)).startswith("module.") | |
| ckpt_has_module = len(ckpt_keys) > 0 and next(iter(ckpt_keys)).startswith("module.") | |
| if model_has_module and not ckpt_has_module: | |
| return {"module." + k: v for k, v in state_dict.items()} | |
| if not model_has_module and ckpt_has_module: | |
| return {k.replace("module.", "", 1): v for k, v in state_dict.items()} | |
| return state_dict | |
| def _is_allowed_missing_key(key): | |
| return key.startswith(ALLOWED_MISSING_PREFIXES) | |
| def load_checkpoint_weights_checked(model, resume_path): | |
| if not resume_path or not os.path.isfile(resume_path): | |
| raise ValueError(f"=> resume failed! no checkpoint found at '{resume_path}'") | |
| logger.info("=> loading checkpoint '{}'".format(resume_path)) | |
| checkpoint = torch.load(resume_path, map_location="cpu") | |
| state_dict = checkpoint.get("state_dict", checkpoint) | |
| state_dict = _normalize_state_dict_for_model(model, state_dict) | |
| incompatible = model.load_state_dict(state_dict, strict=False) | |
| missing = list(incompatible.missing_keys) | |
| unexpected = list(incompatible.unexpected_keys) | |
| bad_missing = [k for k in missing if not _is_allowed_missing_key(k)] | |
| if missing: | |
| shown = missing[:20] | |
| suffix = " ..." if len(missing) > len(shown) else "" | |
| logger.info(f"Missing keys ({len(missing)}): {shown}{suffix}") | |
| if unexpected: | |
| shown = unexpected[:20] | |
| suffix = " ..." if len(unexpected) > len(shown) else "" | |
| logger.info(f"Unexpected keys ({len(unexpected)}): {shown}{suffix}") | |
| if bad_missing or unexpected: | |
| raise RuntimeError( | |
| "Checkpoint/config mismatch for MaPLe inference. " | |
| f"Non-LLaVA missing keys: {bad_missing[:20]} " | |
| f"(total {len(bad_missing)}); unexpected keys: {unexpected[:20]} " | |
| f"(total {len(unexpected)})." | |
| ) | |
| if "epoch" in checkpoint: | |
| logger.info(f"=> loaded checkpoint '{resume_path}' (epoch {checkpoint['epoch']})") | |
| else: | |
| logger.info("=> loaded checkpoint '{}'".format(resume_path)) | |
| return checkpoint | |
| def _ensure_maple_inference_args(args): | |
| network_name = str(args.network_name).lower() | |
| if "maple" not in network_name: | |
| raise ValueError( | |
| f"This custom inference script is for MaPLe checkpoints, got network_name={args.network_name!r}." | |
| ) | |
| if "3dcoord" in network_name or "depthcoord" in network_name: | |
| raise ValueError( | |
| "This script is for ordinary non-depth MaPLe checkpoints. " | |
| "Use a depth-aware custom inference script for maple_3dcoord checkpoints." | |
| ) | |
| if getattr(args, "stage1", False) and not getattr(args, "stage2", False): | |
| raise ValueError( | |
| "The config is still in stage1 mask-only mode. Pass opts like " | |
| "`stage1 False stage2 True use_gt_obj_masks False` for the withft checkpoint." | |
| ) | |
| def _unpack_maple_prediction(pred): | |
| if not isinstance(pred, (tuple, list)) or len(pred) < 5: | |
| raise RuntimeError( | |
| "Expected MaPLe stage2 inference to return five maps " | |
| "(mask, qua, sin, cos, wid). Did you pass `stage1 False stage2 True`?" | |
| ) | |
| return pred[0], pred[1], pred[2], pred[3], pred[4] | |
| def main(): | |
| args = get_parser() | |
| _ensure_maple_inference_args(args) | |
| args.output_dir = os.path.join(args.output_folder, args.exp_name) | |
| setup_logger(args.output_dir, distributed_rank=0, filename="custom_infer_maple.log", mode="a") | |
| logger.info(args) | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA is required for this inference script.") | |
| orig_img, image, inv_mat = preprocess_image(args.custom_image, args.input_size) | |
| h, w = orig_img.shape[:2] | |
| model, _ = build_net(args) | |
| model = torch.nn.DataParallel(model).cuda() | |
| load_checkpoint_weights_checked(model, args.resume) | |
| model.eval() | |
| image = image.cuda(non_blocking=True) | |
| text = [args.custom_sentence] | |
| with torch.no_grad(), amp.autocast(): | |
| pred, _ = model(image, text) | |
| ( | |
| ins_mask_logit, | |
| grasp_qua_logit, | |
| grasp_sin_logit, | |
| grasp_cos_logit, | |
| grasp_wid_logit, | |
| ) = _unpack_maple_prediction(pred) | |
| ins_mask_pred = torch.sigmoid(ins_mask_logit) | |
| grasp_qua_mask_pred = torch.sigmoid(grasp_qua_logit) | |
| grasp_sin_mask_pred = grasp_sin_logit | |
| grasp_cos_mask_pred = grasp_cos_logit | |
| grasp_wid_mask_pred = torch.sigmoid(grasp_wid_logit) | |
| if ins_mask_pred.shape[-2:] != image.shape[-2:]: | |
| pred_size = image.shape[-2:] | |
| ins_mask_pred = F.interpolate( | |
| ins_mask_pred, size=pred_size, mode="bicubic", align_corners=True | |
| ).squeeze(1) | |
| grasp_qua_mask_pred = F.interpolate( | |
| grasp_qua_mask_pred, size=pred_size, mode="bicubic", align_corners=True | |
| ).squeeze(1) | |
| grasp_sin_mask_pred = F.interpolate( | |
| grasp_sin_mask_pred, size=pred_size, mode="bicubic", align_corners=True | |
| ).squeeze(1) | |
| grasp_cos_mask_pred = F.interpolate( | |
| grasp_cos_mask_pred, size=pred_size, mode="bicubic", align_corners=True | |
| ).squeeze(1) | |
| grasp_wid_mask_pred = F.interpolate( | |
| grasp_wid_mask_pred, size=pred_size, mode="bicubic", align_corners=True | |
| ).squeeze(1) | |
| else: | |
| ins_mask_pred = ins_mask_pred.squeeze(1) | |
| grasp_qua_mask_pred = grasp_qua_mask_pred.squeeze(1) | |
| grasp_sin_mask_pred = grasp_sin_mask_pred.squeeze(1) | |
| grasp_cos_mask_pred = grasp_cos_mask_pred.squeeze(1) | |
| grasp_wid_mask_pred = grasp_wid_mask_pred.squeeze(1) | |
| pred_mask_prob = inverse_affine(ins_mask_pred[0].cpu().numpy(), inv_mat, w, h) | |
| pred_mask = pred_mask_prob > args.custom_mask_threshold | |
| grasp_qua_np = inverse_affine(grasp_qua_mask_pred[0].cpu().numpy(), inv_mat, w, h) | |
| grasp_sin_np = inverse_affine(grasp_sin_mask_pred[0].cpu().numpy(), inv_mat, w, h) | |
| grasp_cos_np = inverse_affine(grasp_cos_mask_pred[0].cpu().numpy(), inv_mat, w, h) | |
| grasp_wid_np = inverse_affine(grasp_wid_mask_pred[0].cpu().numpy(), inv_mat, w, h) | |
| logger.info( | |
| "Pred stats: mask_logit min={:.4f}, max={:.4f}, mean={:.4f}; " | |
| "mask_prob min={:.4f}, max={:.4f}, mean={:.4f}, area@{:.2f}={:.4f}; " | |
| "qua min={:.4f}, max={:.4f}, mean={:.4f}; wid min={:.4f}, max={:.4f}, mean={:.4f}; " | |
| "sin mean={:.4f}, cos mean={:.4f}".format( | |
| float(ins_mask_logit.min().detach().cpu()), | |
| float(ins_mask_logit.max().detach().cpu()), | |
| float(ins_mask_logit.mean().detach().cpu()), | |
| float(np.nanmin(pred_mask_prob)), | |
| float(np.nanmax(pred_mask_prob)), | |
| float(np.nanmean(pred_mask_prob)), | |
| float(args.custom_mask_threshold), | |
| float(np.nanmean(pred_mask.astype(np.float32))), | |
| float(np.nanmin(grasp_qua_np)), | |
| float(np.nanmax(grasp_qua_np)), | |
| float(np.nanmean(grasp_qua_np)), | |
| float(np.nanmin(grasp_wid_np)), | |
| float(np.nanmax(grasp_wid_np)), | |
| float(np.nanmean(grasp_wid_np)), | |
| float(np.nanmean(grasp_sin_np)), | |
| float(np.nanmean(grasp_cos_np)), | |
| ) | |
| ) | |
| grasp_qua_for_detect = grasp_qua_np | |
| if getattr(args, "custom_gate_grasps_by_mask", False): | |
| if pred_mask.any(): | |
| grasp_qua_for_detect = grasp_qua_np * pred_mask.astype(grasp_qua_np.dtype) | |
| logger.info("Using predicted mask to gate final grasp peak detection.") | |
| else: | |
| grasp_qua_for_detect = np.zeros_like(grasp_qua_np) | |
| logger.warning( | |
| "Predicted mask is empty at the current threshold; " | |
| "mask-gated grasp detection will return no peaks." | |
| ) | |
| pred_grasps, _ = detect_grasps( | |
| grasp_qua_for_detect, | |
| grasp_sin_np, | |
| grasp_cos_np, | |
| grasp_wid_np, | |
| args.custom_num_grasps, | |
| ) | |
| if pred_grasps: | |
| logger.info(f"Pred grasps [cx, cy, width, height, theta_deg]: {pred_grasps}") | |
| else: | |
| logger.warning("No grasp peaks found by detect_grasps.") | |
| save_path = args.custom_save or default_save_path(args) | |
| save_prediction_overlay(orig_img, pred_mask, pred_grasps, args.custom_sentence, save_path) | |
| logger.info(f"Saved MaPLe custom inference visualization: {save_path}") | |
| print(save_path) | |
| if __name__ == "__main__": | |
| main() | |