Untitled

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'
)