diff --git a/tools/inference.txt b/tools/inference.txt new file mode 100644 index 0000000000..c67518f2fa --- /dev/null +++ b/tools/inference.txt @@ -0,0 +1,14 @@ +python tools/inference_custom.py \ + /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/td-hm_res50_8xb64-210e_coco-256x192.py \ + /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/epoch_300.pth \ + --img-dir /data/new_mmpose/mmpose/data/1704_split_exported_data_project_id_422/val2017 \ + --out-dir /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/out \ + --bbox-json /data/new_mmpose/mmpose/data/1704_split_exported_data_project_id_422/annotations/forklift_keypoints_val2017.json \ + --output-file /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/results.json + +python tools/inference_custom.py \ + /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/td-hm_res50_8xb64-210e_coco-256x192.py \ + /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/epoch_300.pth \ + --img-dir /data/new_mmpose/mmpose/data/1704_split_exported_data_project_id_422/val2017 \ + --bbox-json /data/new_mmpose/mmpose/data/1704_split_exported_data_project_id_422/annotations/forklift_keypoints_val2017.json \ + --output-file /data/new_mmpose/mmpose/work_dirs/ls_1704_res18/results.json \ No newline at end of file diff --git a/tools/inference_custom.py b/tools/inference_custom.py new file mode 100644 index 0000000000..a05edcd81d --- /dev/null +++ b/tools/inference_custom.py @@ -0,0 +1,131 @@ +import argparse +import os +import json +import numpy as np +import cv2 +from mmengine.config import Config, DictAction +from mmpose.apis import init_model, inference_topdown +from xtcocotools.coco import COCO +from mmengine.utils import ProgressBar + + +def parse_args(): + parser = argparse.ArgumentParser(description="Run MMPose inference on images") + parser.add_argument('config', help='Path to model config file') + parser.add_argument('model', help='Path to checkpoint or ONNX file') + parser.add_argument('--img-dir', type=str, required=True, help='Directory with input images') + parser.add_argument('--bbox-json', type=str, required=True, help='Path to COCO format bounding box JSON') + parser.add_argument('--out-dir', type=str, help='Directory to save visualized results (optional)') + parser.add_argument('--output-file', type=str, help='File to save keypoint results in JSON') + parser.add_argument('--device', default='cuda:0', help='Device to run inference on (e.g., "cuda:0" or "cpu")') + parser.add_argument('--score-thr', type=float, default=0.3, help='Keypoint score threshold') + parser.add_argument( + '--cfg-options', nargs='+', action=DictAction, + help='Override some settings in the config file. The key-value pair in ' + 'xxx=yyy format will be merged into the config.') + return parser.parse_args() + + +def draw_keypoints(image, keypoints, scores, score_thr): + """Draw keypoints on the image using OpenCV.""" + for el1, el2 in zip(keypoints, scores): + for kp, score in zip(el1, el2): + if score > score_thr: + x, y = int(kp[0]), int(kp[1]) + cv2.circle(image, (x, y), 5, (0, 255, 0), -1) # Draw keypoint + cv2.putText(image, f"{int(score*100)}%", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) + + +def main(): + args = parse_args() + + # Load the COCO bounding boxes + coco = COCO(args.bbox_json) + img_ids = list(coco.imgs.keys()) + + # Initialize the model + cfg = Config.fromfile(args.config) + if args.cfg_options: + cfg.merge_from_dict(args.cfg_options) + model = init_model(cfg, args.model, device=args.device) + + # Ensure output directories exist if `out-dir` is provided + if args.out_dir: + os.makedirs(args.out_dir, exist_ok=True) + + # Results to be saved + results = [] + + # Progress bar + progress_bar = ProgressBar(len(img_ids)) + + for img_id in img_ids: + img_info = coco.loadImgs([img_id])[0] + img_path = os.path.join(args.img_dir, img_info['file_name']) + + if not os.path.exists(img_path): + print(f"Image not found: {img_path}") + progress_bar.update() + continue + + # Load the image + image = cv2.imread(img_path) + if image is None: + print(f"Failed to read image: {img_path}") + progress_bar.update() + continue + + # Convert image to RGB for inference + image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + # Load bounding boxes for the image + ann_ids = coco.getAnnIds(imgIds=[img_id]) + annotations = coco.loadAnns(ann_ids) + person_bboxes = np.array([ann['bbox'] for ann in annotations]) + + # Run pose inference + pose_results = inference_topdown( + model, + image_rgb, + person_bboxes, + bbox_format='xywh' # COCO annotations typically use 'xywh' format + ) + + # Extract keypoints and bounding boxes from PoseDataSample + keypoints_results = [] + for pose in pose_results: + pred_instances = pose.pred_instances + if pred_instances is not None: + keypoints = pred_instances.keypoints + scores = pred_instances.keypoint_scores + keypoints_results.append({ + 'keypoints': keypoints.tolist(), + 'scores': scores.tolist() + }) + + # Draw keypoints on the image + draw_keypoints(image, keypoints, scores, args.score_thr) + + # Save the visualized image if `out-dir` is provided + if args.out_dir: + out_file = os.path.join(args.out_dir, img_info['file_name']) + cv2.imwrite(out_file, image) + + # Save keypoints to results + results.append({ + 'image_id': img_id, + 'file_name': img_info['file_name'], + 'keypoints': keypoints_results + }) + + progress_bar.update() + + # Save results to output file + if args.output_file: + with open(args.output_file, 'w') as f: + json.dump(results, f, indent=4) + print(f"Inference completed. Results {'saved to ' + args.out_dir if args.out_dir else ''}.") + + +if __name__ == '__main__': + main()