Untitled
unknown
plain_text
5 months ago
3.6 kB
5
Indexable
from mmpose.apis import init_pose_model, inference_top_down_pose_model, vis_pose_result import os import cv2 import matplotlib.pyplot as plt def inference_and_render_skeleton(image_path, output_path): # Paths config_file = 'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/ochuman/ViTPose_base_ochuman_256x192.py' checkpoint_file = '/content/drive/MyDrive/ViTPose/models/vitpose_base_coco_aic_mpii.pth' image_folder = image_path output_folder = output_path os.makedirs(output_folder, exist_ok=True) # Initialize the model pose_model = init_pose_model(config_file, checkpoint_file, device='cpu') # cuda:0 line_thickness = 3 keypoint_radius = 4 output_data = [] for image_name in os.listdir(image_folder): print(image_name) image_path = os.path.join(image_folder, image_name) if os.path.isdir(image_path): continue # Read image image = cv2.imread(image_path) height, width, _ = image.shape print(width, height) # Dummy bounding box for the whole image (if no person detection) dummy_bboxes = [{'bbox': [0, 0, image.shape[1], image.shape[0], 1]}] # Perform inference pose_results, _ = inference_top_down_pose_model( pose_model, image_path, person_results=dummy_bboxes, format='xywh', bbox_thr=0.5, dataset='TopDownOCHumanDataset', # Set for OCHuman return_heatmap=False ) vis_result = vis_pose_result( pose_model, image_path, pose_results, dataset='TopDownOCHumanDataset', show=False, kpt_score_thr=0.5, # Confidence threshold for keypoints radius=keypoint_radius, thickness=line_thickness ) vis_result_rgb = cv2.cvtColor(vis_result, cv2.COLOR_BGR2RGB) # Show the image in the notebook plt.imshow(vis_result_rgb) plt.axis('off') plt.show() cv2.imwrite(os.path.join(output_folder, image_name), vis_result) print(f"Processed {image_name}") for result in pose_results: # Loop over all detected results result['keypoints'] = result['keypoints'].astype(np.float64) keypoints = [] bbox = [float('inf'), float('inf'), -float('inf'), -float('inf')] for kp in result['keypoints']: # Keypoints is a list of (x, y, visibility) x, y, visibility = kp keypoints.extend([x, y, visibility]) if visibility >= 0.5: bbox[0] = min(bbox[0], x) bbox[1] = min(bbox[1], y) bbox[2] = max(bbox[2], x) bbox[3] = max(bbox[3], y) center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2] scale = [bbox[2] - bbox[0], bbox[3] - bbox[1]] # center, scale, score doesn't affect mAP calculation. Use dummy values output_entry = { "category_id": 1, "center": [0, 0], "image_id": name2id.get(image_name, 1), "keypoints": keypoints, "scale": [0, 0], "score": 1.0 } output_data.append(output_entry) with open(os.path.join(output_folder, 'annotations.json'), 'w') as f: json.dump(output_data, f, indent=4) print(f"Annotations saved")
Editor is loading...
Leave a Comment