Untitled
unknown
plain_text
5 months ago
2.8 kB
3
Indexable
import cv2 import numpy as np import matplotlib.pyplot as plt import json # Function to visualize predictions on a single image def render_skeleton_from_predictions(image_path, image_id, result_path, save_path): with open(result_path, 'r') as f: results = json.load(f) bbox_color = (255, 0, 0) # Red for bounding boxes skeleton_color = (0, 255, 0) # Green for skeletons # Define skeleton connections (based on COCO format) SKELETON = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] for i in range(len(SKELETON)): SKELETON[i][0] -= 1 SKELETON[i][1] -= 1 image = cv2.imread(image_path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Filter results for the given image_id filtered_results = [result for result in results if result['image_id'] == image_id] if not filtered_results: print(f"No predictions found for image_id {image_id}") return for result in filtered_results: print(result) keypoints = np.array(result['keypoints']).reshape(-1, 3) # x, y, confidence bbox_center = result['center'] bbox_scale = result['scale'] image_width, image_height = image.shape[1], image.shape[0] print(image_width, image_height) # Draw keypoints for kp in keypoints: if kp[2] > 0.5: # Confidence threshold cv2.circle(image, (int(kp[0]), int(kp[1])), 5, skeleton_color, -1) # Draw skeleton for start, end in SKELETON: if start < len(keypoints) and end < len(keypoints): if keypoints[start][2] > 0.5 and keypoints[end][2] > 0.5: # Check confidence for both points cv2.line(image, (int(keypoints[start][0]), int(keypoints[start][1])), (int(keypoints[end][0]), int(keypoints[end][1])), skeleton_color, 2) # Show the image plt.figure(figsize=(10, 10)) plt.imshow(image) plt.axis('off') #plt.title(f"Predictions for {image_path}") # Save image plt.savefig(save_path, bbox_inches='tight', pad_inches=0) print(f"Image saved to {save_path}") plt.show() # Find the corresponding image_id from the JSON file #image_id = name2id.get(image_name, 1) render_skeleton_from_predictions( '/content/drive/MyDrive/ViTPose/inference/000008/original/000008.jpg', 1, '/content/drive/MyDrive/ViTPose/inference/000008/inpainted/output/annotations.json', '/content/drive/MyDrive/ViTPose/inference/000008/inpainted/output/predictions.jpg' )
Editor is loading...
Leave a Comment